PointMLP/classification_ModelNet40/main.py
2021-10-04 04:00:01 -04:00

250 lines
11 KiB
Python

import argparse
import os
import logging
import datetime
import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
from torch.utils.data import DataLoader
import models as models
from utils import Logger, mkdir_p, progress_bar, save_model, save_args, cal_loss
from data import ModelNet40
from torch.optim.lr_scheduler import CosineAnnealingLR
import sklearn.metrics as metrics
import numpy as np
def parse_args():
"""Parameters"""
parser = argparse.ArgumentParser('training')
parser.add_argument('-c', '--checkpoint', type=str, metavar='PATH',
help='path to save checkpoint (default: checkpoint)')
parser.add_argument('--msg', type=str, help='message after checkpoint')
parser.add_argument('--batch_size', type=int, default=32, help='batch size in training')
parser.add_argument('--model', default='PointNet', help='model name [default: pointnet_cls]')
parser.add_argument('--epoch', default=300, type=int, help='number of epoch in training')
parser.add_argument('--num_points', type=int, default=1024, help='Point Number')
parser.add_argument('--learning_rate', default=0.1, type=float, help='learning rate in training')
parser.add_argument('--weight_decay', type=float, default=2e-4, help='decay rate')
parser.add_argument('--seed', type=int, help='random seed')
parser.add_argument('--workers', default=8, type=int, help='workers')
return parser.parse_args()
def main():
args = parse_args()
if args.seed is None:
args.seed = np.random.randint(1, 10000)
os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
assert torch.cuda.is_available(), "Please ensure codes are executed in cuda."
device = 'cuda'
if args.seed is not None:
torch.manual_seed(args.seed)
np.random.seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
torch.cuda.manual_seed(args.seed)
torch.set_printoptions(10)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
os.environ['PYTHONHASHSEED'] = str(args.seed)
time_str = str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S'))
if args.msg is None:
message = time_str
else:
message = "-" + args.msg
args.checkpoint = 'checkpoints/' + args.model + message + '-' + str(args.seed)
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
screen_logger = logging.getLogger("Model")
screen_logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(message)s')
file_handler = logging.FileHandler(os.path.join(args.checkpoint, "out.txt"))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
screen_logger.addHandler(file_handler)
def printf(str):
screen_logger.info(str)
print(str)
# Model
printf(f"args: {args}")
printf('==> Building model..')
net = models.__dict__[args.model]()
criterion = cal_loss
net = net.to(device)
# criterion = criterion.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
best_test_acc = 0. # best test accuracy
best_train_acc = 0.
best_test_acc_avg = 0.
best_train_acc_avg = 0.
best_test_loss = float("inf")
best_train_loss = float("inf")
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
optimizer_dict = None
if not os.path.isfile(os.path.join(args.checkpoint, "last_checkpoint.pth")):
save_args(args)
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title="ModelNet" + args.model)
logger.set_names(["Epoch-Num", 'Learning-Rate',
'Train-Loss', 'Train-acc-B', 'Train-acc',
'Valid-Loss', 'Valid-acc-B', 'Valid-acc'])
else:
printf(f"Resuming last checkpoint from {args.checkpoint}")
checkpoint_path = os.path.join(args.checkpoint, "last_checkpoint.pth")
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint['net'])
start_epoch = checkpoint['epoch']
best_test_acc = checkpoint['best_test_acc']
best_train_acc = checkpoint['best_train_acc']
best_test_acc_avg = checkpoint['best_test_acc_avg']
best_train_acc_avg = checkpoint['best_train_acc_avg']
best_test_loss = checkpoint['best_test_loss']
best_train_loss = checkpoint['best_train_loss']
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title="ModelNet" + args.model, resume=True)
optimizer_dict = checkpoint['optimizer']
printf('==> Preparing data..')
train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=args.workers,
batch_size=args.batch_size, shuffle=True, drop_last=True)
test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=args.workers,
batch_size=args.batch_size // 2, shuffle=False, drop_last=False)
optimizer = torch.optim.SGD(net.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay)
if optimizer_dict is not None:
optimizer.load_state_dict(optimizer_dict)
scheduler = CosineAnnealingLR(optimizer, args.epoch, eta_min=1e-3, last_epoch=start_epoch - 1)
for epoch in range(start_epoch, args.epoch):
printf('Epoch(%d/%s) Learning Rate %s:' % (epoch + 1, args.epoch, optimizer.param_groups[0]['lr']))
train_out = train(net, train_loader, optimizer, criterion, device) # {"loss", "acc", "acc_avg", "time"}
test_out = validate(net, test_loader, criterion, device)
scheduler.step()
if test_out["acc"] > best_test_acc:
best_test_acc = test_out["acc"]
is_best = True
else:
is_best = False
best_test_acc = test_out["acc"] if (test_out["acc"] > best_test_acc) else best_test_acc
best_train_acc = train_out["acc"] if (train_out["acc"] > best_train_acc) else best_train_acc
best_test_acc_avg = test_out["acc_avg"] if (test_out["acc_avg"] > best_test_acc_avg) else best_test_acc_avg
best_train_acc_avg = train_out["acc_avg"] if (train_out["acc_avg"] > best_train_acc_avg) else best_train_acc_avg
best_test_loss = test_out["loss"] if (test_out["loss"] < best_test_loss) else best_test_loss
best_train_loss = train_out["loss"] if (train_out["loss"] < best_train_loss) else best_train_loss
save_model(
net, epoch, path=args.checkpoint, acc=test_out["acc"], is_best=is_best,
best_test_acc=best_test_acc, # best test accuracy
best_train_acc=best_train_acc,
best_test_acc_avg=best_test_acc_avg,
best_train_acc_avg=best_train_acc_avg,
best_test_loss=best_test_loss,
best_train_loss=best_train_loss,
optimizer=optimizer.state_dict()
)
logger.append([epoch, optimizer.param_groups[0]['lr'],
train_out["loss"], train_out["acc_avg"], train_out["acc"],
test_out["loss"], test_out["acc_avg"], test_out["acc"]])
printf(
f"Training loss:{train_out['loss']} acc_avg:{train_out['acc_avg']}% acc:{train_out['acc']}% time:{train_out['time']}s")
printf(
f"Testing loss:{test_out['loss']} acc_avg:{test_out['acc_avg']}% "
f"acc:{test_out['acc']}% time:{test_out['time']}s [best test acc: {best_test_acc}%] \n\n")
logger.close()
printf(f"++++++++" * 2 + "Final results" + "++++++++" * 2)
printf(f"++ Last Train time: {train_out['time']} | Last Test time: {test_out['time']} ++")
printf(f"++ Best Train loss: {best_train_loss} | Best Test loss: {best_test_loss} ++")
printf(f"++ Best Train acc_B: {best_train_acc_avg} | Best Test acc_B: {best_test_acc_avg} ++")
printf(f"++ Best Train acc: {best_train_acc} | Best Test acc: {best_test_acc} ++")
printf(f"++++++++" * 5)
def train(net, trainloader, optimizer, criterion, device):
net.train()
train_loss = 0
correct = 0
total = 0
train_pred = []
train_true = []
time_cost = datetime.datetime.now()
for batch_idx, (data, label) in enumerate(trainloader):
data, label = data.to(device), label.to(device).squeeze()
data = data.permute(0, 2, 1) # so, the input data shape is [batch, 3, 1024]
optimizer.zero_grad()
logits = net(data)
loss = criterion(logits, label)
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), 1)
optimizer.step()
train_loss += loss.item()
preds = logits.max(dim=1)[1]
train_true.append(label.cpu().numpy())
train_pred.append(preds.detach().cpu().numpy())
total += label.size(0)
correct += preds.eq(label).sum().item()
progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (train_loss / (batch_idx + 1), 100. * correct / total, correct, total))
time_cost = int((datetime.datetime.now() - time_cost).total_seconds())
train_true = np.concatenate(train_true)
train_pred = np.concatenate(train_pred)
return {
"loss": float("%.3f" % (train_loss / (batch_idx + 1))),
"acc": float("%.3f" % (100. * metrics.accuracy_score(train_true, train_pred))),
"acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(train_true, train_pred))),
"time": time_cost
}
def validate(net, testloader, criterion, device):
net.eval()
test_loss = 0
correct = 0
total = 0
test_true = []
test_pred = []
time_cost = datetime.datetime.now()
with torch.no_grad():
for batch_idx, (data, label) in enumerate(testloader):
data, label = data.to(device), label.to(device).squeeze()
data = data.permute(0, 2, 1)
logits = net(data)
loss = criterion(logits, label)
test_loss += loss.item()
preds = logits.max(dim=1)[1]
test_true.append(label.cpu().numpy())
test_pred.append(preds.detach().cpu().numpy())
total += label.size(0)
correct += preds.eq(label).sum().item()
progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (test_loss / (batch_idx + 1), 100. * correct / total, correct, total))
time_cost = int((datetime.datetime.now() - time_cost).total_seconds())
test_true = np.concatenate(test_true)
test_pred = np.concatenate(test_pred)
return {
"loss": float("%.3f" % (test_loss / (batch_idx + 1))),
"acc": float("%.3f" % (100. * metrics.accuracy_score(test_true, test_pred))),
"acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(test_true, test_pred))),
"time": time_cost
}
if __name__ == '__main__':
main()