class PointNetTrainer(): def __init__(self): self.model = PointNet().to(device) #self.model = PFFNet().to(device) self.lr = 1e-4 self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr) self.criterion = torch.nn.CrossEntropyLoss() self.epoches = 20 def train(self, loader): self.model.train() total_loss = 0 tot_num = 0 for i, data in enumerate(tqdm.tqdm(loader)): tot_num += len(data.y) data = data.to(device) self.optimizer.zero_grad() logits = self.model(data.pos, data.batch) loss = self.criterion(logits, data.y) loss.backward() self.optimizer.step() total_loss += loss.item() return total_loss / tot_num @torch.no_grad() def test(self, loader): self.model.eval() total_correct = 0 tot_num = 0 for i, data in enumerate(loader): if i > 50: break tot_num += len(data.y) data = data.to(device) logits = self.model(data.pos, data.batch) pred = logits.argmax(dim=-1) total_correct += (pred == data.y).sum() #print(total_correct,tot_num) return total_correct / tot_num def work(self, train_loader, test_loader): plt.figure() losses = [] accs = [] for epoch in range(self.epoches): loss = self.train(train_loader) acc = self.test(test_loader) print("epoch", epoch, "loss", loss, "acc", acc) losses.append(loss) accs.append(acc) plt.plot(losses) plt.plot(accs) plt.legend(['loss', 'acc']) plt.savefig('dump/curve.png')
def test(args, io): if args.dataset == 'modelnet40': test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) elif args.dataset == 'ScanObjectNN': test_loader = DataLoader(ScanObjectNN(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) else: raise Exception("Dataset Not supported") device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': if args.dataset == 'modelnet40': model = PointNet(args, output_channels=40).to(device) elif args.dataset == 'ScanObjectNN': model = PointNet(args, output_channels=15).to(device) else: raise Exception("Dataset Not supported") elif args.model == 'dgcnn': if args.dataset == 'modelnet40': model = DGCNN(args, output_channels=40).to(device) elif args.dataset == 'ScanObjectNN': model = DGCNN(args, output_channels=15).to(device) else: raise Exception("Dataset Not supported") elif args.model == 'gbnet': if args.dataset == 'modelnet40': model = GBNet(args, output_channels=40).to(device) elif args.dataset == 'ScanObjectNN': model = GBNet(args, output_channels=15).to(device) else: raise Exception("Dataset Not supported") else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) 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] 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), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] SHAPE_NAMES = [line.rstrip() for line in \ open('data/modelnet40_ply_hdf5_2048/shape_names.txt')] NUM_CLASSES = 40 total_seen_class = [0 for _ in range(NUM_CLASSES)] total_correct_class = [0 for _ in range(NUM_CLASSES)] 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) 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) for i in range(test_true.shape[0]): l = test_true[i] total_seen_class[l] += 1 total_correct_class[l] += (test_pred[i] == l) class_accuracies = np.array(total_correct_class) / np.array( total_seen_class, dtype=np.float) for i, name in enumerate(SHAPE_NAMES): io.cprint('%10s:\t%0.3f' % (name, class_accuracies[i]))
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN_cls(args).to(device) else: raise Exception("Not implemented") model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) 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] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) # visualize - added by jaeha if args.visualize: xyz = data[0].cpu() ax = plt.axes(projection='3d') ax.scatter(xyz[0, :], xyz[1, :], xyz[2, :], s=1, color='blue') plt.title('True: ' + class_lists[label[0]] + ' , Pred: ' + class_lists[preds[0]]) plt.show() 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), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN_cls(args).to(device) else: raise Exception("Not implemented") model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) # model.load_state_dict(torch.load("/home/mask/xas_ws/dgcnn_pytorch/checkpoints/cls_1024/models/model.t7")) 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] 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 mytest(args, io): # test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), # batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN_cls(args).to(device) else: raise Exception("Not implemented") model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] pc_test = pc_from_h5('/home/mask/MSCNN/data-generation/resources/dataset/object_4e/real_test/1.h5') # for i in range(10): pc_test = torch.from_numpy(pc_test) pc_test = pc_test.type(torch.FloatTensor) pc_test = pc_test.cuda() pc_test = pc_test.permute(0, 2, 1) pc_test = pc_test.to(device) logits = model(pc_test) preds = logits.max(dim=1)[1] test_pred.append(preds.detach().cpu().numpy()) print(test_pred)
state = torch.load(ckp_path) model.load_state_dict(state['state_dict']) print("model load from %s" % ckp_path) if __name__ == "__main__": torch.manual_seed(SEED) device = torch.device(f'cuda:{gpus[0]}' if torch.cuda.is_available() else 'cpu') print("Loading test dataset...") test_data = PointNetDataset("./dataset/modelnet40_normal_resampled", train=1) test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True) model = PointNet().to(device=device) if ckp_path: load_ckp(ckp_path, model) model = model.to(device) model.eval() with torch.no_grad(): accs = [] gt_ys = [] pred_ys = [] for x, y in test_loader: x = x.to(device) y = y.to(device) # TODO: put x into network and get out out = # TODO: get pred_y from out pred_y =
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) 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) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) io.cprint(outstr) #################### # 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) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name)
def startcustomtraining(args, io): ft_loader = DataLoader(FT10(num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=True) ft_test_loader = DataLoader(FT11(num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) criterion = cal_loss best_ft_test_acc = 0.0 i = 0 train_accs = [] test_accs = [] epochs = [] for epoch in range(args.epochs): i += 1 scheduler.step() ft_loss = 0.0 count = 0 model.train() ft_pred = [] ft_true = [] for data, label in ft_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 ft_loss += loss.item() * batch_size ft_true.append(label.cpu().numpy()) ft_pred.append(preds.detach().cpu().numpy()) #print(data.shape, label.shape, logits.shape, preds.shape) #print('LABELS:', label) #print('PREDS:', preds) #print('LOGITS:', logits) ft_true = np.concatenate(ft_true) ft_pred = np.concatenate(ft_pred) ft_acc = metrics.accuracy_score(ft_true, ft_pred) avg_per_class_acc = metrics.balanced_accuracy_score(ft_true, ft_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, ft_loss * 1.0 / count, ft_acc, avg_per_class_acc) io.cprint(outstr) train_accs.append(ft_acc) ft_test_loss = 0.0 count = 0 model.eval() ft_test_pred = [] ft_test_true = [] for data, label in ft_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 ft_test_loss += loss.item() * batch_size ft_test_true.append(label.cpu().numpy()) ft_test_pred.append(preds.detach().cpu().numpy()) #print(data.shape, label.shape, logits.shape, preds.shape) #print('LABELS:', label) #print('PREDS:', preds) #print('LOGITS:', logits) ft_test_true = np.concatenate(ft_test_true) ft_test_pred = np.concatenate(ft_test_pred) ft_test_acc = metrics.accuracy_score(ft_test_true, ft_test_pred) avg_per_class_acc = metrics.balanced_accuracy_score( ft_test_true, ft_test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, ft_test_loss * 1.0 / count, ft_test_acc, avg_per_class_acc) io.cprint(outstr) if ft_test_acc > best_ft_test_acc: print('save now') best_ft_test_acc = ft_test_acc torch.save(model.state_dict(), 'pretrained/custommodel.t7') #torch.save(model.state_dict(), 'pretrained/custommodel.t7') epochs.append(i) test_accs.append(ft_test_acc) fig, ax = plt.subplots() ax.plot(epochs, train_accs, color='blue', label='train acc') ax.plot(epochs, test_accs, color='red', label='test acc') ax.set(xlabel='epoch', ylabel='accuracy', title='accuracy values per epoch') ax.grid() ax.legend() fig.savefig("accuracy.png") plt.show()
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) elif args.model == 'semigcn': model = SemiGCN(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) opt.load_state_dict(checkpoint['opt']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) #scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr, last_epoch=args.start_epoch-1) scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.8)#0.7 #scheduler = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.9825, last_epoch=args.start_epoch-1) criterion = cal_loss best_test_acc = 0 for epoch in range(args.start_epoch, 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()) scheduler.step() train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % (epoch, train_loss*1.0/count, metrics.accuracy_score( train_true, train_pred), metrics.balanced_accuracy_score( train_true, train_pred)) io.cprint(outstr) if epoch%10 == 0: # save running checkpoint per 10 epoch torch.save({'epoch': epoch + 1, 'arch': args.model, 'state_dict': model.state_dict(), 'opt' : opt.state_dict()}, 'checkpoints/%s/models/checkpoint_latest.pth.tar' % args.exp_name) #################### # 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) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % (epoch, test_loss*1.0/count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save({'epoch': epoch + 1, 'arch': args.model, 'state_dict': model.state_dict(), 'opt' : opt.state_dict()}, 'checkpoints/%s/models/checkpoint_best.pth.tar' % args.exp_name)
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) elif args.model == 'ssg': model = PointNet2SSG(output_classes=40, dropout_prob=args.dropout) model.to(device) elif args.model == 'msg': model = PointNet2MSG(output_classes=40, dropout_prob=args.dropout) model.to(device) elif args.model == 'ognet': # [64,128,256,512] model = Model_dense(20, args.feature_dims, [512], output_classes=40, init_points=768, input_dims=3, dropout_prob=args.dropout, id_skip=args.id_skip, drop_connect_rate=args.drop_connect_rate, cluster='xyzrgb', pre_act=args.pre_act, norm=args.norm_layer) if args.efficient: model = ModelE_dense(20, args.feature_dims, [512], output_classes=40, init_points=768, input_dims=3, dropout_prob=args.dropout, id_skip=args.id_skip, drop_connect_rate=args.drop_connect_rate, cluster='xyzrgb', pre_act=args.pre_act, norm=args.norm_layer, gem=args.gem, ASPP=args.ASPP) model.to(device) elif args.model == 'ognet-small': # [48,96,192,384] model = Model_dense(20, args.feature_dims, [512], output_classes=40, init_points=768, input_dims=3, dropout_prob=args.dropout, id_skip=args.id_skip, drop_connect_rate=args.drop_connect_rate, cluster='xyzrgb', pre_act=args.pre_act, norm=args.norm_layer) model.to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=0.01 * args.lr) criterion = cal_loss best_test_acc = 0 best_avg_per_class_acc = 0 warm_up = 0.1 # We start from the 0.1*lrRate warm_iteration = round( len(ModelNet40(partition='train', num_points=args.num_points)) / args.batch_size) * args.warm_epoch # first 5 epoch 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() batch_size = data.size()[0] opt.zero_grad() if args.model == 'ognet' or args.model == 'ognet-small' or args.model == 'ssg' or args.model == 'msg': logits = model(data, data) else: data = data.permute(0, 2, 1) logits = model(data) loss = criterion(logits, label) if epoch < args.warm_epoch: warm_up = min(1.0, warm_up + 0.9 / warm_iteration) loss *= warm_up 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) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) io.cprint(outstr) #################### # 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() batch_size = data.size()[0] if args.model == 'ognet' or args.model == 'ognet-small' or args.model == 'ssg' or args.model == 'msg': logits = model(data, data) else: data = data.permute(0, 2, 1) 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) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc + avg_per_class_acc >= best_test_acc + best_avg_per_class_acc: best_test_acc = test_acc best_avg_per_class_acc = avg_per_class_acc print('This is the current best.') torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) elif args.model == 'ssg': model = PointNet2SSG(output_classes=40, dropout_prob=0) model.to(device) elif args.model == 'msg': model = PointNet2MSG(output_classes=40, dropout_prob=0) model.to(device) elif args.model == 'ognet': # [64,128,256,512] model = Model_dense(20, args.feature_dims, [512], output_classes=40, init_points=768, input_dims=3, dropout_prob=args.dropout, id_skip=args.id_skip, drop_connect_rate=args.drop_connect_rate, cluster='xyzrgb', pre_act=args.pre_act, norm=args.norm_layer) if args.efficient: model = ModelE_dense(20, args.feature_dims, [512], output_classes=40, init_points=768, input_dims=3, dropout_prob=args.dropout, id_skip=args.id_skip, drop_connect_rate=args.drop_connect_rate, cluster='xyzrgb', pre_act=args.pre_act, norm=args.norm_layer, gem=args.gem, ASPP=args.ASPP) model.to(device) elif args.model == 'ognet-small': # [48,96,192,384] model = Model_dense(20, args.feature_dims, [512], output_classes=40, init_points=768, input_dims=3, dropout_prob=args.dropout, id_skip=args.id_skip, drop_connect_rate=args.drop_connect_rate, cluster='xyzrgb', pre_act=args.pre_act, norm=args.norm_layer) model.to(device) else: raise Exception("Not implemented") try: model.load_state_dict(torch.load(args.model_path)) except: model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() model = model.module batch0, label0 = next(iter(test_loader)) batch0 = batch0[0].unsqueeze(0) print(batch0.shape) print(model) macs, params = get_model_complexity_info(model, batch0, ((1024, 3)), as_strings=True, print_per_layer_stat=False, verbose=True) print('{:<30} {:<8}'.format('Computational complexity: ', macs)) print('{:<30} {:<8}'.format('Number of parameters: ', params)) 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() batch_size = data.size()[0] if args.model == 'ognet' or args.model == 'ognet-small' or args.model == 'ssg' or args.model == 'msg': logits = model(data, data) #logits = model(1.1*data, 1.1*data) else: data = data.permute(0, 2, 1) 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(modelin=args.model,outfile=args.out,optimize=args.opt): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model','calib_' + modelin) sfm_path = os.path.join('model','sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() # mean shape and eigenvectors for 3dmm M = 100 data3dmm = dataloader.SyntheticLoader() mu_lm = torch.from_numpy(data3dmm.mu_lm).float().detach() mu_lm[:,2] = mu_lm[:,2]*-1 lm_eigenvec = torch.from_numpy(data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) # sample from f testing set allerror_2d = [] allerror_3d = [] allerror_rel3d = [] allerror_relf = [] all_f = [] all_fpred = [] all_depth = [] out_shape = [] out_f = [] seterror_3d = [] seterror_rel3d = [] seterror_relf = [] seterror_2d = [] f_vals = [i*100 for i in range(4,15)] for f_test in f_vals: # create dataloader loader = dataloader.TestLoader(f_test) f_pred = [] shape_pred = [] error_2d = [] error_3d = [] error_rel3d = [] error_relf = [] M = 100; N = 68; batch_size = 1; for j,data in enumerate(loader): if j == 10: break # load the data x_cam_gt = data['x_cam_gt'] shape_gt = data['x_w_gt'] fgt = data['f_gt'] x_img = data['x_img'] x_img_gt = data['x_img_gt'] T_gt = data['T_gt'] all_depth.append(np.mean(T_gt[:,2])) all_f.append(fgt.numpy()[0]) ptsI = x_img.reshape((M,N,2)).permute(0,2,1) x = x_img.unsqueeze(0).permute(0,2,1) # run the model f = calib_net(x) + 300 betas = sfm_net(x) betas = betas.squeeze(0).unsqueeze(-1) shape = mu_lm + torch.mm(lm_eigenvec,betas).squeeze().view(N,3) # additional optimization on initial solution if optimize: calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.train() sfm_net.train() opt1 = torch.optim.Adam(calib_net.parameters(),lr=1e-4) opt2 = torch.optim.Adam(sfm_net.parameters(),lr=1e-2) curloss = 100 for outerloop in itertools.count(): # camera calibration shape = shape.detach() for iter in itertools.count(): opt1.zero_grad() print(x.shape) quit() f = calib_net(x) + 300 K = torch.zeros(3,3).float() K[0,0] = f K[1,1] = f K[2,2] = 1 f_error = torch.mean(torch.abs(f - fgt)) rmse = torch.norm(shape_gt - shape,dim=1).mean() # differentiable PnP pose estimation km,c_w,scaled_betas, alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI,K) error2d = util.getReprojError2(ptsI,shape,R,T,K,show=False,loss='l1') #error2d = util.getReprojError2_(ptsI,Xc,K,show=True,loss='l1') error_time = util.getTimeConsistency(shape,R,T) loss = error2d.mean() + 0.01*error_time if iter == 5: break loss.backward() opt1.step() print(f"iter: {iter} | error: {loss.item():.3f} | f/fgt: {f.item():.1f}/{fgt[0].item():.1f} | error2d: {error2d.mean().item():.3f} | rmse: {rmse.item():.3f} ") # sfm f = f.detach() for iter in itertools.count(): opt2.zero_grad() # shape prediction betas = sfm_net(x) shape = torch.sum(betas * lm_eigenvec,1) shape = shape.reshape(68,3) + mu_lm K = torch.zeros((3,3)).float() K[0,0] = f K[1,1] = f K[2,2] = 1 #rmse = torch.norm(shape_gt - shape,dim=1).mean().detach() rmse = torch.norm(shape_gt - shape,dim=1).mean().detach() # differentiable PnP pose estimation km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI,K) error2d = util.getReprojError2(ptsI,shape,R,T,K,show=False,loss='l1') #loss = rmse loss = error2d.mean() if iter == 5: break if iter > 10 and prev_loss < loss: break else: prev_loss = loss loss.backward() opt2.step() print(f"iter: {iter} | error: {loss.item():.3f} | f/fgt: {f.item():.1f}/{fgt[0].item():.1f} | error2d: {error2d.mean().item():.3f} | rmse: {rmse.item():.3f} ") # closing condition for outerloop on dual objective if torch.abs(curloss - loss) < 0.01: break curloss = loss else: K = torch.zeros(3,3).float() K[0,0] = f K[1,1] = f K[2,2] = 1 km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI,K) all_fpred.append(f.detach().numpy()[0]) # get errors reproj_errors2 = util.getReprojError2(ptsI,shape,R,T,K,show=False) reproj_errors3 = torch.norm(shape_gt - shape,dim=1).mean() rel_errors = util.getRelReprojError3(x_cam_gt,shape,R,T) reproj_error = reproj_errors2.mean() reconstruction_error = reproj_errors3.mean() rel_error = rel_errors.mean() f_error = torch.abs(fgt - f) / fgt # save final prediction f_pred.append(f.detach().cpu().item()) shape_pred.append(shape.detach().cpu().numpy()) allerror_3d.append(reproj_error.data.numpy()) allerror_2d.append(reconstruction_error.data.numpy()) allerror_rel3d.append(rel_error.data.numpy()) error_2d.append(reproj_error.cpu().data.item()) error_3d.append(reconstruction_error.cpu().data.item()) error_rel3d.append(rel_error.cpu().data.item()) error_relf.append(f_error.cpu().data.item()) print(f"f/sequence: {f_test}/{j} | f/fgt: {f[0].item():.3f}/{fgt.item():.3f} | f_error_rel: {f_error.item():.4f} | rmse: {reconstruction_error.item():.4f} | rel rmse: {rel_error.item():.4f} | 2d error: {reproj_error.item():.4f}") avg_2d = np.mean(error_2d) avg_rel3d = np.mean(error_rel3d) avg_3d = np.mean(error_3d) avg_relf = np.mean(error_relf) seterror_2d.append(avg_2d) seterror_3d.append(avg_3d) seterror_rel3d.append(avg_rel3d) seterror_relf.append(avg_relf) out_f.append(np.stack(f_pred)) out_shape.append(np.stack(shape_pred,axis=0)) print(f"f_error_rel: {avg_relf:.4f} | rel rmse: {avg_rel3d:.4f} | 2d error: {reproj_error.item():.4f} | rmse: {avg_3d:.4f} |") out_shape = np.stack(out_shape) out_f = np.stack(out_f) all_f = np.stack(all_f).flatten() all_fpred = np.stack(all_fpred).flatten() all_d = np.stack(all_depth).flatten() allerror_2d = np.stack(allerror_2d).flatten() allerror_3d = np.stack(allerror_3d).flatten() allerror_rel3d = np.stack(allerror_rel3d).flatten() matdata = {} matdata['fvals'] = np.array(f_vals) matdata['all_f'] = np.array(all_f) matdata['all_fpred'] = np.array(all_fpred) matdata['all_d'] = np.array(all_depth) matdata['error_2d'] = allerror_2d matdata['error_3d'] = allerror_3d matdata['error_rel3d'] = allerror_rel3d matdata['seterror_2d'] = np.array(seterror_2d) matdata['seterror_3d'] = np.array(seterror_3d) matdata['seterror_rel3d'] = np.array(seterror_rel3d) matdata['seterror_relf'] = np.array(seterror_relf) matdata['shape'] = np.stack(out_shape) matdata['f'] = np.stack(out_f) scipy.io.savemat(outfile,matdata) print(f"MEAN seterror_2d: {np.mean(seterror_2d)}") print(f"MEAN seterror_3d: {np.mean(seterror_3d)}") print(f"MEAN seterror_rel3d: {np.mean(seterror_rel3d)}") print(f"MEAN seterror_relf: {np.mean(seterror_relf)}")
def testBIWIID(modelin=args.model,outfile=args.out,optimize=args.opt): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model','calib_' + modelin) sfm_path = os.path.join('model','sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() # mean shape and eigenvectors for 3dmm data3dmm = dataloader.SyntheticLoader() mu_lm = torch.from_numpy(data3dmm.mu_lm).float().detach() mu_lm[:,2] = mu_lm[:,2]*-1 lm_eigenvec = torch.from_numpy(data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) # define loader loader = dataloader.BIWIIDLoader() f_pred = [] shape_pred = [] error_2d = [] error_relf = [] error_rel3d = [] for idx in range(len(loader)): batch = loader[idx] x_cam_gt = batch['x_cam_gt'] fgt = batch['f_gt'] x_img = batch['x_img'] x_img_gt = batch['x_img_gt'] M = x_img_gt.shape[0] N = 68 ptsI = x_img.reshape((M,N,2)).permute(0,2,1) x = ptsI.unsqueeze(0).permute(0,2,1,3) # run the model f = calib_net(x) + 300 betas = sfm_net(x) betas = betas.squeeze(0).unsqueeze(-1) shape = mu_lm + torch.mm(lm_eigenvec,betas).squeeze().view(N,3) # additional optimization on initial solution if optimize: calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.train() sfm_net.train() opt1 = torch.optim.Adam(calib_net.parameters(),lr=1e-4) opt2 = torch.optim.Adam(sfm_net.parameters(),lr=1e-2) curloss = 100 for outerloop in itertools.count(): # camera calibration shape = shape.detach() for iter in itertools.count(): opt1.zero_grad() f = calib_net.forward2(x) + 300 K = torch.zeros(3,3).float() K[0,0] = f K[1,1] = f K[2,2] = 1 f_error = torch.mean(torch.abs(f - fgt)) #rmse = torch.norm(shape_gt - shape,dim=1).mean() # differentiable PnP pose estimation km,c_w,scaled_betas, alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI,K) error2d = util.getReprojError2(ptsI,shape,R,T,K,show=False,loss='l1') error_time = util.getTimeConsistency(shape,R,T) #error_shape = util.get3DConsistency(ptsI,shape,kinv,R,T) order = torch.pow(10,-1*torch.floor(torch.log10(error_time)).detach()) #loss = error2d.mean() + order*error_time loss = error2d.mean() if iter == 5: break #if iter > 10 and prev_loss < loss: # break #else: # prev_loss = loss loss.backward() opt1.step() print(f"iter: {iter} | error: {loss.item():.3f} | f/fgt: {f.item():.1f}/{fgt[0].item():.1f} | error2d: {error2d.mean().item():.3f} ") # sfm f = f.detach() for iter in itertools.count(): opt2.zero_grad() # shape prediction betas = sfm_net.forward2(x) shape = torch.sum(betas * lm_eigenvec,1) shape = shape.reshape(68,3) + mu_lm K = torch.zeros((3,3)).float() K[0,0] = f K[1,1] = f K[2,2] = 1 #rmse = torch.norm(shape_gt - shape,dim=1).mean().detach() #rmse = torch.norm(shape_gt - shape,dim=1).mean().detach() # differentiable PnP pose estimation km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI,K) error2d = util.getReprojError2(ptsI,shape,R,T,K,show=False,loss='l1') #loss = rmse loss = error2d.mean() if iter == 5: break if iter > 10 and prev_loss < loss: break else: prev_loss = loss loss.backward() opt2.step() print(f"iter: {iter} | error: {loss.item():.3f} | f/fgt: {f.item():.1f}/{fgt[0].item():.1f} | error2d: {error2d.mean().item():.3f} ") # closing condition for outerloop on dual objective if torch.abs(curloss - loss) < 0.01: break curloss = loss else: K = torch.zeros(3,3).float() K[0,0] = f K[1,1] = f K[2,2] = 1 km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI,K) # get errors reproj_errors2 = util.getReprojError2(ptsI,shape,R,T,K) rel_errors = util.getRelReprojError3(x_cam_gt,shape,R,T) reproj_error = reproj_errors2.mean() rel_error = rel_errors.mean() f_error = torch.abs(fgt - f) / fgt # save final prediction f_pred.append(f.detach().cpu().item()) shape_pred.append(shape.detach().cpu().numpy()) error_2d.append(reproj_error.cpu().data.item()) error_rel3d.append(rel_error.cpu().data.item()) error_relf.append(f_error.cpu().data.item()) print(f" f/fgt: {f[0].item():.3f}/{fgt.item():.3f} | f_error_rel: {f_error.item():.4f} | rel rmse: {rel_error.item():.4f} | 2d error: {reproj_error.item():.4f}") #end for # prepare output file out_shape = np.stack(shape_pred) out_f = np.stack(f_pred) matdata = {} matdata['shape'] = np.stack(out_shape) matdata['f'] = np.stack(out_f) matdata['error_2d'] = np.array(error_2d) matdata['error_rel3d'] = np.array(error_rel3d) matdata['error_relf'] = np.array(error_relf) scipy.io.savemat(outfile,matdata) print(f"MEAN seterror_2d: {np.mean(error_2d)}") print(f"MEAN seterror_rel3d: {np.mean(error_rel3d)}") print(f"MEAN seterror_relf: {np.mean(error_relf)}")
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train'), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test'), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=False) device = torch.device("cuda:0") model = PointNet().to(device) print(str(model)) print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-6) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) 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() batch_size = data.size()[0] opt.zero_grad() logits = model(data.float()) 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) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) io.cprint(outstr) #################### # 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() batch_size = data.size()[0] logits = model(data.float()) 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) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name) print('Saving ckpt with acc: %f' % best_test_acc)
class FID(object): def __init__(self, mode, dataset, device, split, path=None): if mode == "PointNet": self.model = PointNet().to(device) if path is None: path = "./metrics/pointnet_modelnet40/checkpoints/pointnet_max_pc_2048_emb_1024/models/model.t7" if split == 'train': self.real_stat_save_path = "./metrics/gt_stats/pointnet_max_pc_2048_emb_1024/%s-train.npz" % dataset elif split == 'test': self.real_stat_save_path = "./metrics/gt_stats/pointnet_max_pc_2048_emb_1024/%s-test.npz" % dataset else: raise ValueError('ERROR: unknown split %s!' % split) print('Using PointNet, gt_stat_fn: %s\n' % self.real_stat_save_path) else: raise ValueError('ERROR: unknown FID mode %s!' % mode) self.model.load_state_dict(torch.load(path)) self.model = self.model.eval() self.device = device def get_fid(self, fake_pts, batch_size=32): f = np.load(self.real_stat_save_path) real_mean, real_cov = f['mean'], f['cov'] fake_feature_list = [] with torch.no_grad(): b, _, _ = fake_pts.shape for i in range(b // batch_size + 1): fake_pts_batch = fake_pts[i * batch_size:min((i + 1) * batch_size, b)] if fake_pts_batch.shape[0] > 0: fake_feature_batch = self.model( torch.Tensor(fake_pts_batch).to( self.device))[1].cpu().detach().numpy() fake_feature_list.append(fake_feature_batch) fake_feature = np.concatenate(fake_feature_list, 0) fake_mean = np.mean(fake_feature, axis=0) fake_cov = np.cov(fake_feature, rowvar=False) fid = self.calculate_frechet_distance(real_mean, real_cov, fake_mean, fake_cov) return fid @staticmethod def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): """Numpy implementation of the Frechet Distance. The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) and X_2 ~ N(mu_2, C_2) is d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)). Stable version by Dougal J. Sutherland. Params: -- mu1 : Numpy array containing the activations of a layer of the inception net (like returned by the function 'get_predictions') for generated samples. -- mu2 : The sample mean over activations, precalculated on an representative data set. -- sigma1: The covariance matrix over activations for generated samples. -- sigma2: The covariance matrix over activations, precalculated on an representative data set. Returns: -- : The Frechet Distance. """ mu1 = np.atleast_1d(mu1) mu2 = np.atleast_1d(mu2) sigma1 = np.atleast_2d(sigma1) sigma2 = np.atleast_2d(sigma2) assert mu1.shape == mu2.shape, \ 'Training and test mean vectors have different lengths' assert sigma1.shape == sigma2.shape, \ 'Training and test covariances have different dimensions' diff = mu1 - mu2 # Product might be almost singular covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) if not np.isfinite(covmean).all(): msg = ('fid calculation produces singular product; ' 'adding %s to diagonal of cov estimates') % eps print(msg) offset = np.eye(sigma1.shape[0]) * eps covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) # Numerical error might give slight imaginary component if np.iscomplexobj(covmean): if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): m = np.max(np.abs(covmean.imag)) raise ValueError('Imaginary component {}'.format(m)) covmean = covmean.real tr_covmean = np.trace(covmean) return (diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)