def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points, pt_norm=False), batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") # Try to load models: if args.arch == 'dgcnn': from model.DGCNN_PAConv import PAConv model = PAConv(args).to(device) elif args.arch == 'pointnet': from model.PointNet_PAConv import PAConv model = PAConv(args).to(device) else: raise Exception("Not implemented") io.cprint(str(model)) model = nn.DataParallel(model) model.load_state_dict( torch.load("checkpoints/%s/best_model.t7" % args.exp_name)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] with torch.no_grad(): logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points, pt_norm=False), num_workers=args.workers, batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") NUM_PEPEAT = 300 NUM_VOTE = 10 # Try to load models: if args.arch == 'dgcnn': from model.DGCNN_PAConv import PAConv model = PAConv(args).to(device) elif args.arch == 'pointnet': from model.PointNet_PAConv import PAConv model = PAConv(args).to(device) else: raise Exception("Not implemented") model = nn.DataParallel(model) model.load_state_dict( torch.load("checkpoints/%s/best_model.t7" % args.exp_name)) model = model.eval() best_acc = 0 pointscale = PointcloudScale(scale_low=0.8, scale_high=1.18) # set the range of scaling for i in range(NUM_PEPEAT): test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() pred = 0 for v in range(NUM_VOTE): new_data = data batch_size = data.size()[0] if v > 0: new_data.data = pointscale(new_data.data) with torch.no_grad(): pred += F.softmax(model(new_data.permute(0, 2, 1)), dim=1) # sum 10 preds pred /= NUM_VOTE # avg the preds! label = label.view(-1) pred_choice = pred.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(pred_choice.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) if test_acc > best_acc: best_acc = test_acc outstr = 'Voting %d, test acc: %.6f,' % (i, test_acc * 100) io.cprint(outstr) final_outstr = 'Final voting test acc: %.6f,' % (best_acc * 100) io.cprint(final_outstr)
def train(args, io): # ============= Model =================== num_part = 50 device = torch.device("cuda" if args.cuda else "cpu") model = PAConv(args, num_part).to(device) io.cprint(str(model)) model.apply(weight_init) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") '''Use Pretrain or not''' if args.get('pretrain', False): state_dict = torch.load("checkpoints/%s/best_insiou_model.pth" % args.exp_name, map_location=torch.device('cpu'))['model'] for k in state_dict.keys(): if 'module' not in k: from collections import OrderedDict new_state_dict = OrderedDict() for k in state_dict: new_state_dict['module.' + k] = state_dict[k] state_dict = new_state_dict break model.load_state_dict(state_dict) print("Using pretrained model...") print(torch.load("checkpoints/%s/best_insiou_model.pth" % args.exp_name).keys()) else: print("Training from scratch...") # =========== Dataloader ================= train_data = PartNormalDataset(npoints=2048, split='trainval', normalize=False) print("The number of training data is:%d", len(train_data)) test_data = PartNormalDataset(npoints=2048, split='test', normalize=False) print("The number of test data is:%d", len(test_data)) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, drop_last=True) test_loader = DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, drop_last=False) # ============= Optimizer ================ if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.weight_decay) if args.scheduler == 'cos': print("Use CosLR") scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr / 100) else: print("Use StepLR") scheduler = StepLR(opt, step_size=args.step, gamma=0.5) # ============= Training ================= best_acc = 0 best_class_iou = 0 best_instance_iou = 0 num_part = 50 num_classes = 16 for epoch in range(args.epochs): train_epoch(train_loader, model, opt, scheduler, epoch, num_part, num_classes, io) test_metrics, total_per_cat_iou = test_epoch(test_loader, model, epoch, num_part, num_classes, io) # 1. when get the best accuracy, save the model: if test_metrics['accuracy'] > best_acc: best_acc = test_metrics['accuracy'] io.cprint('Max Acc:%.5f' % best_acc) state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_acc': best_acc} torch.save(state, 'checkpoints/%s/best_acc_model.pth' % args.exp_name) # 2. when get the best instance_iou, save the model: if test_metrics['shape_avg_iou'] > best_instance_iou: best_instance_iou = test_metrics['shape_avg_iou'] io.cprint('Max instance iou:%.5f' % best_instance_iou) state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_instance_iou': best_instance_iou} torch.save(state, 'checkpoints/%s/best_insiou_model.pth' % args.exp_name) # 3. when get the best class_iou, save the model: # first we need to calculate the average per-class iou class_iou = 0 for cat_idx in range(16): class_iou += total_per_cat_iou[cat_idx] avg_class_iou = class_iou / 16 if avg_class_iou > best_class_iou: best_class_iou = avg_class_iou # print the iou of each class: for cat_idx in range(16): io.cprint(classes_str[cat_idx] + ' iou: ' + str(total_per_cat_iou[cat_idx])) io.cprint('Max class iou:%.5f' % best_class_iou) state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_class_iou': best_class_iou} torch.save(state, 'checkpoints/%s/best_clsiou_model.pth' % args.exp_name) # report best acc, ins_iou, cls_iou io.cprint('Final Max Acc:%.5f' % best_acc) io.cprint('Final Max instance iou:%.5f' % best_instance_iou) io.cprint('Final Max class iou:%.5f' % best_class_iou) # save last model state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': args.epochs - 1, 'test_iou': best_instance_iou} torch.save(state, 'checkpoints/%s/model_ep%d.pth' % (args.exp_name, args.epochs))
def test(args, io): # Dataloader test_data = PartNormalDataset(npoints=2048, split='test', normalize=False) print("The number of test data is:%d", len(test_data)) test_loader = DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, drop_last=False) # Try to load models num_part = 50 device = torch.device("cuda" if args.cuda else "cpu") model = PAConv(args, num_part).to(device) io.cprint(str(model)) from collections import OrderedDict state_dict = torch.load("checkpoints/%s/best_%s_model.pth" % (args.exp_name, args.model_type), map_location=torch.device('cpu'))['model'] new_state_dict = OrderedDict() for layer in state_dict: new_state_dict[layer.replace('module.', '')] = state_dict[layer] model.load_state_dict(new_state_dict) model.eval() num_part = 50 num_classes = 16 metrics = defaultdict(lambda: list()) hist_acc = [] shape_ious = [] total_per_cat_iou = np.zeros((16)).astype(np.float32) total_per_cat_seen = np.zeros((16)).astype(np.int32) for batch_id, (points, label, target, norm_plt) in tqdm(enumerate(test_loader), total=len(test_loader), smoothing=0.9): batch_size, num_point, _ = points.size() points, label, target, norm_plt = Variable(points.float()), Variable(label.long()), Variable(target.long()), Variable(norm_plt.float()) points = points.transpose(2, 1) norm_plt = norm_plt.transpose(2, 1) points, label, target, norm_plt = points.cuda(non_blocking=True), label.squeeze().cuda( non_blocking=True), target.cuda(non_blocking=True), norm_plt.cuda(non_blocking=True) with torch.no_grad(): seg_pred = model(points, norm_plt, to_categorical(label, num_classes)) # b,n,50 # instance iou without considering the class average at each batch_size: batch_shapeious = compute_overall_iou(seg_pred, target, num_part) # [b] shape_ious += batch_shapeious # iou +=, equals to .append # per category iou at each batch_size: for shape_idx in range(seg_pred.size(0)): # sample_idx cur_gt_label = label[shape_idx] # label[sample_idx] total_per_cat_iou[cur_gt_label] += batch_shapeious[shape_idx] total_per_cat_seen[cur_gt_label] += 1 # accuracy: seg_pred = seg_pred.contiguous().view(-1, num_part) target = target.view(-1, 1)[:, 0] pred_choice = seg_pred.data.max(1)[1] correct = pred_choice.eq(target.data).cpu().sum() metrics['accuracy'].append(correct.item() / (batch_size * num_point)) hist_acc += metrics['accuracy'] metrics['accuracy'] = np.mean(hist_acc) metrics['shape_avg_iou'] = np.mean(shape_ious) for cat_idx in range(16): if total_per_cat_seen[cat_idx] > 0: total_per_cat_iou[cat_idx] = total_per_cat_iou[cat_idx] / total_per_cat_seen[cat_idx] # First we need to calculate the iou of each class and the avg class iou: class_iou = 0 for cat_idx in range(16): class_iou += total_per_cat_iou[cat_idx] io.cprint(classes_str[cat_idx] + ' iou: ' + str(total_per_cat_iou[cat_idx])) # print the iou of each class avg_class_iou = class_iou / 16 outstr = 'Test :: test acc: %f test class mIOU: %f, test instance mIOU: %f' % (metrics['accuracy'], avg_class_iou, metrics['shape_avg_iou']) io.cprint(outstr)
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points, pt_norm=args.pt_norm), 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, pt_norm=False), num_workers=args.workers, batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") if args.arch == 'dgcnn': from model.DGCNN_PAConv import PAConv model = PAConv(args).to(device) elif args.arch == 'pointnet': from model.PointNet_PAConv import PAConv model = PAConv(args).to(device) else: raise Exception("Not implemented") io.cprint(str(model)) model.apply(weight_init) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr / 100) criterion = cal_loss best_test_acc = 0 for epoch in range(args.epochs): scheduler.step() #################### # Train #################### train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] for data, label in train_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() logits = model(data) loss = criterion(logits, label) loss.backward() opt.step() preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) train_acc = metrics.accuracy_score(train_true, train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, ' % ( epoch, train_loss * 1.0 / count, train_acc) io.cprint(outstr) writer.add_scalar('loss_train', train_loss * 1.0 / count, epoch + 1) writer.add_scalar('Acc_train', train_acc, epoch + 1) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) loss = criterion(logits, label) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f,' % ( epoch, test_loss * 1.0 / count, test_acc) io.cprint(outstr) writer.add_scalar('loss_test', test_loss * 1.0 / count, epoch + 1) writer.add_scalar('Acc_test', test_acc, epoch + 1) if test_acc >= best_test_acc: best_test_acc = test_acc io.cprint('Max Acc:%.6f' % best_test_acc) torch.save(model.state_dict(), 'checkpoints/%s/best_model.t7' % args.exp_name)
def train(gpu, ngpus_per_node): # ============= Model =================== if args.arch == 'dgcnn': from model.DGCNN_PAConv import PAConv model = PAConv(args) elif args.arch == 'pointnet': from model.PointNet_PAConv import PAConv model = PAConv(args) else: raise Exception("Not implemented") model.apply(weight_init) if main_process(): logger.info(model) if args.sync_bn and args.distributed: model = nn.SyncBatchNorm.convert_sync_batchnorm(model) if args.distributed: torch.cuda.set_device(gpu) args.batch_size = int(args.batch_size / ngpus_per_node) args.test_batch_size = int(args.test_batch_size / ngpus_per_node) args.workers = int( (args.workers + ngpus_per_node - 1) / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel( model.cuda(), device_ids=[gpu], find_unused_parameters=True) else: model = torch.nn.DataParallel(model.cuda()) # =========== Dataloader ================= train_data = ModelNet40(partition='train', num_points=args.num_points, pt_norm=args.pt_norm) test_data = ModelNet40(partition='test', num_points=args.num_points, pt_norm=False) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_data) test_sampler = torch.utils.data.distributed.DistributedSampler( test_data) else: train_sampler = None test_sampler = None train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True) test_loader = torch.utils.data.DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=test_sampler) # ============= Optimizer =================== if main_process(): logger.info("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr / 100) criterion = cal_loss best_test_acc = 0 start_epoch = 0 # ============= Training from scratch================= for epoch in range(start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) train_epoch(train_loader, model, opt, scheduler, epoch, criterion) test_acc = test_epoch(test_loader, model, epoch, criterion) if test_acc >= best_test_acc and main_process(): best_test_acc = test_acc logger.info('Max Acc:%.6f' % best_test_acc) torch.save(model.state_dict(), 'checkpoints/%s/best_model.t7' % args.exp_name) # save the best model
def test(gpu, ngpus_per_node): if main_process(): logger.info('<<<<<<<<<<<<<<<<< Start Evaluation <<<<<<<<<<<<<<<<<') # ============= Model =================== if args.arch == 'dgcnn': from model.DGCNN_PAConv import PAConv model = PAConv(args) elif args.arch == 'pointnet': from model.PointNet_PAConv import PAConv model = PAConv(args) else: raise Exception("Not implemented") if main_process(): logger.info(model) if args.sync_bn: assert args.distributed == True model = nn.SyncBatchNorm.convert_sync_batchnorm(model) if args.distributed: torch.cuda.set_device(gpu) args.batch_size = int(args.batch_size / ngpus_per_node) args.test_batch_size = int(args.test_batch_size / ngpus_per_node) args.workers = int( (args.workers + ngpus_per_node - 1) / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel( model.cuda(), device_ids=[gpu], find_unused_parameters=True) else: model = torch.nn.DataParallel(model.cuda()) state_dict = torch.load("checkpoints/%s/best_model.t7" % args.exp_name, map_location=torch.device('cpu')) for k in state_dict.keys(): if 'module' not in k: from collections import OrderedDict new_state_dict = OrderedDict() for k in state_dict: new_state_dict['module.' + k] = state_dict[k] state_dict = new_state_dict break model.load_state_dict(state_dict) # Dataloader test_data = ModelNet40(partition='test', num_points=args.num_points) if args.distributed: test_sampler = torch.utils.data.distributed.DistributedSampler( test_data) else: test_sampler = None test_loader = torch.utils.data.DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=test_sampler) model.eval() intersection_meter = AverageMeter() union_meter = AverageMeter() target_meter = AverageMeter() for data, label in test_loader: data, label = data.cuda(non_blocking=True), label.cuda( non_blocking=True).squeeze(1) data = data.permute(0, 2, 1) with torch.no_grad(): logits = model(data) preds = logits.max(dim=1)[1] intersection, union, target = intersectionAndUnionGPU( preds, label, args.classes) if args.multiprocessing_distributed: dist.all_reduce(intersection), dist.all_reduce( union), dist.all_reduce(target) intersection, union, target = intersection.cpu().numpy(), union.cpu( ).numpy(), target.cpu().numpy() intersection_meter.update(intersection), union_meter.update( union), target_meter.update(target) accuracy_class = intersection_meter.sum / (target_meter.sum + 1e-10) mAcc = np.mean(accuracy_class) allAcc = sum(intersection_meter.sum) / (sum(target_meter.sum) + 1e-10) if main_process(): logger.info('Test result: mAcc/allAcc {:.4f}/{:.4f}.'.format( mAcc, allAcc)) for i in range(args.classes): logger.info('Class_{} Result: accuracy {:.4f}.'.format( i, accuracy_class[i])) logger.info('<<<<<<<<<<<<<<<<< End Evaluation <<<<<<<<<<<<<<<<<')
def train(gpu, ngpus_per_node): # ============= Model =================== num_part = 50 model = PAConv(args, num_part) model.apply(weight_init) if main_process(): logger.info(model) if args.sync_bn and args.distributed: model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) if args.distributed: torch.cuda.set_device(gpu) args.batch_size = int(args.batch_size / ngpus_per_node) args.test_batch_size = int(args.test_batch_size / ngpus_per_node) args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[gpu], find_unused_parameters=True) else: model = torch.nn.DataParallel(model.cuda()) '''Use Pretrain or not''' if args.get('pretrain', False): state_dict = torch.load("checkpoints/%s/best_insiou_model.pth" % args.exp_name, map_location=torch.device('cpu'))['model'] for k in state_dict.keys(): if 'module' not in k: from collections import OrderedDict new_state_dict = OrderedDict() for k in state_dict: new_state_dict['module.' + k] = state_dict[k] state_dict = new_state_dict break model.load_state_dict(state_dict) if main_process(): logger.info("Using pretrained model...") logger.info(torch.load("checkpoints/%s/best_insiou_model.pth" % args.exp_name).keys()) else: if main_process(): logger.info("Training from scratch...") # =========== Dataloader ================= train_data = PartNormalDataset(npoints=2048, split='trainval', normalize=False) if main_process(): logger.info("The number of training data is:%d", len(train_data)) test_data = PartNormalDataset(npoints=2048, split='test', normalize=False) if main_process(): logger.info("The number of test data is:%d", len(test_data)) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_data) test_sampler = torch.utils.data.distributed.DistributedSampler(test_data) else: train_sampler = None test_sampler = None train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) test_loader = torch.utils.data.DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=test_sampler) # ============= Optimizer =================== if args.use_sgd: if main_process(): logger.info("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) else: if main_process(): logger.info("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.weight_decay) if args.scheduler == 'cos': if main_process(): logger.info("Use CosLR") scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr / 100) else: if main_process(): logger.info("Use StepLR") scheduler = StepLR(opt, step_size=args.step, gamma=0.5) # ============= Training ================= best_acc = 0 best_class_iou = 0 best_instance_iou = 0 num_part = 50 num_classes = 16 for epoch in range(args.epochs): if args.distributed: train_sampler.set_epoch(epoch) train_epoch(train_loader, model, opt, scheduler, epoch, num_part, num_classes) test_metrics, total_per_cat_iou = test_epoch(test_loader, model, epoch, num_part, num_classes) # 1. when get the best accuracy, save the model: if test_metrics['accuracy'] > best_acc and main_process(): best_acc = test_metrics['accuracy'] logger.info('Max Acc:%.5f' % best_acc) state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_acc': best_acc} torch.save(state, 'checkpoints/%s/best_acc_model.pth' % args.exp_name) # 2. when get the best instance_iou, save the model: if test_metrics['shape_avg_iou'] > best_instance_iou and main_process(): best_instance_iou = test_metrics['shape_avg_iou'] logger.info('Max instance iou:%.5f' % best_instance_iou) state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_instance_iou': best_instance_iou} torch.save(state, 'checkpoints/%s/best_insiou_model.pth' % args.exp_name) # 3. when get the best class_iou, save the model: # first we need to calculate the average per-class iou class_iou = 0 for cat_idx in range(16): class_iou += total_per_cat_iou[cat_idx] avg_class_iou = class_iou / 16 if avg_class_iou > best_class_iou and main_process(): best_class_iou = avg_class_iou # print the iou of each class: for cat_idx in range(16): if main_process(): logger.info(classes_str[cat_idx] + ' iou: ' + str(total_per_cat_iou[cat_idx])) logger.info('Max class iou:%.5f' % best_class_iou) state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_class_iou': best_class_iou} torch.save(state, 'checkpoints/%s/best_clsiou_model.pth' % args.exp_name) if main_process(): # report best acc, ins_iou, cls_iou logger.info('Final Max Acc:%.5f' % best_acc) logger.info('Final Max instance iou:%.5f' % best_instance_iou) logger.info('Final Max class iou:%.5f' % best_class_iou) # save last model state = { 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), 'optimizer': opt.state_dict(), 'epoch': args.epochs - 1, 'test_iou': best_instance_iou} torch.save(state, 'checkpoints/%s/model_ep%d.pth' % (args.exp_name, args.epochs))