def main():
print('Loss & Optimizer')
if args.loss == 'triplet':
args.triplet = True
criterion = TripletLoss(margin=args.margin, swap=args.swap)
elif args.loss == 'triplet_distance':
args.triplet = True
criterion = TripletLoss(margin=args.margin, swap=args.swap, dist=True)
else:
args.triplet = False
criterion = ContrastiveLoss(margin=args.margin)
print('Prepare data')
train_loader, valid_loader, valid_gallery_loader, test_loader, test_gallery_loader, in_size = load_data(
args.dataset,
args.data_path,
triplet=args.triplet,
batch_size=args.batch_size,
prefetch=args.prefetch)
print('Create model')
net = models.GNN(in_size,
args.out_size,
nlayers=args.nlayers,
hid=args.hidden,
J=args.pow)
distNet = distance.SoftHd()
print('Check CUDA')
if args.cuda and args.ngpu > 1:
print('\t* Data Parallel **NOT TESTED**')
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
if args.cuda:
print('\t* CUDA')
net, distNet = net.cuda(), distNet.cuda()
criterion = criterion.cuda()
start_epoch = 0
best_map = 0
early_stop_counter = 0
if args.load is not None:
print('Loading model')
checkpoint = load_checkpoint(args.load)
net.load_state_dict(checkpoint['state_dict'])
distNet.load_state_dict(checkpoint['state_dict_dist'])
start_epoch = checkpoint['epoch']
best_map = checkpoint['best_map']
print('Loaded model at epoch {epoch} and mAP {meanap}%'.format(
epoch=checkpoint['epoch'], meanap=checkpoint['best_map']))
print('***Test***')
test(test_loader, test_gallery_loader, [net, distNet], args.cuda)
def main(query, query_name, target, target_name):
print('Prepare data')
split = os.path.normpath(args.data_path).split(os.sep)
split[-2] = split[-2] + '-pickled'
pickle_dir = os.path.join(*split)
if split[0]=='':
pickle_dir = os.sep + pickle_dir
# gt_path = os.path.join(args.data_path, os.pardir, '00_GroundTruth', args.set_partition)
gt_path = os.path.join(args.data_path, os.pardir, '00_GroundTruth')
# data = HistoGraph(pickle_dir, os.path.join(gt_path, 'test.txt'))
data = HistoGraph(os.path.join(pickle_dir, '02_Test'), os.path.join(gt_path, '02_Test', 'words.txt'))
data_query = HistoGraph(os.path.join(pickle_dir, '02_Test'), os.path.join(gt_path, '02_Test', 'queries.txt'))
# data_query = copy.deepcopy(data)
data_query.graphs = [query]
data_query.labels = [query_name]
g1, l1 = data_query[0]
data_target = data
data_target.graphs = [target]
data_target.labels = [target_name]
g2, l2 = data_target[0]
del data_target, data_query, data
print('Create model')
net = models.GNN(2, args.hidden, args.out_size, dropout=args.dropout)
distNet = distance.SoftHd(args.out_size)
print('Check CUDA')
if args.cuda:
print('\t* CUDA')
net, distNet = net.cuda(), distNet.cuda()
g1.ndata['pos'] = g1.ndata['pos'].cuda()
g2.ndata['pos'] = g2.ndata['pos'].cuda()
if args.load is not None:
print('Loading model')
checkpoint = load_checkpoint(args.load)
net.load_state_dict(checkpoint['state_dict'])
distNet.load_state_dict(checkpoint['state_dict_dist'])
print('***PLOT***')
g1 = net(g1)
g2 = net(g2)
dist, indB, indA = distNet.soft_hausdorff(g1, g2, train=False)
plot(g1.to_networkx(node_attrs=['pos']).to_undirected(), g2.to_networkx(node_attrs=['pos']).to_undirected(), indB.tolist(), indA.tolist(), query_name, target_name, dist.item())
def main():
print('Loss & Optimizer')
if args.loss == 'triplet':
args.triplet = True
criterion = TripletLoss(margin=args.margin, swap=args.swap)
elif args.loss == 'triplet_distance':
args.triplet = True
criterion = TripletLoss(margin=args.margin, swap=args.swap, dist=True)
else:
args.triplet = False
criterion = ContrastiveLoss(margin=args.margin)
print('Prepare data')
train_loader, valid_loader, valid_gallery_loader, test_loader, test_gallery_loader, in_size = load_data(
args.dataset,
args.data_path,
triplet=args.triplet,
batch_size=args.batch_size,
prefetch=args.prefetch,
set_partition=args.set_partition)
print('Create model')
if args.model == 'GAT':
net = models.GNN_GAT(in_size,
args.hidden,
args.out_size,
dropout=args.dropout)
elif args.model == 'GRU':
net = models.GNN_GRU(in_size,
args.hidden,
args.out_size,
dropout=args.dropout)
distNet = distance.SoftHd(args.out_size)
optimizer = torch.optim.Adam(list(net.parameters()) +
list(distNet.parameters()),
args.learning_rate,
weight_decay=args.decay)
scheduler = StepLR(optimizer, 5, gamma=args.gamma)
print('Check CUDA')
if args.cuda and args.ngpu > 1:
print('\t* Data Parallel **NOT TESTED**')
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
if args.cuda:
print('\t* CUDA')
net, distNet = net.cuda(), distNet.cuda()
criterion = criterion.cuda()
start_epoch = 0
best_perf = 0
early_stop_counter = 0
if args.load is not None:
print('Loading model')
checkpoint = load_checkpoint(args.load)
net.load_state_dict(checkpoint['state_dict'])
distNet.load_state_dict(checkpoint['state_dict_dist'])
start_epoch = checkpoint['epoch']
best_perf = checkpoint['best_perf']
if not args.test:
print('***Train***')
for epoch in range(start_epoch, args.epochs):
loss_train = train(train_loader, [net, distNet], optimizer,
args.cuda, criterion, epoch)
acc_valid, map_valid = test(valid_loader,
valid_gallery_loader, [net, distNet],
args.cuda,
validation=True)
# Early-Stop + Save model
if map_valid.avg > best_perf:
best_perf = map_valid.avg
early_stop_counter = 0
if args.save is not None:
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': net.state_dict(),
'state_dict_dist': distNet.state_dict(),
'best_perf': best_perf
},
directory=args.save,
file_name='checkpoint')
else:
if early_stop_counter >= args.early_stop:
print('Early Stop epoch {}'.format(epoch))
break
early_stop_counter += 1
# Logger
if args.log:
# Scalars
logger.add_scalar('loss_train', loss_train.avg)
logger.add_scalar('acc_valid', acc_valid.avg)
logger.add_scalar('map_valid', map_valid.avg)
logger.add_scalar('learning_rate', scheduler.get_lr()[0])
logger.step()
scheduler.step()
# Load Best model in case of save it
if args.save is not None:
print('Loading best model')
best_model_file = os.path.join(args.save, 'checkpoint.pth')
checkpoint = load_checkpoint(best_model_file)
net.load_state_dict(checkpoint['state_dict'])
distNet.load_state_dict(checkpoint['state_dict_dist'])
print('Best model at epoch {epoch} and acc {acc}%'.format(
epoch=checkpoint['epoch'], acc=checkpoint['best_perf']))
print('***Valid***')
test(valid_loader, valid_gallery_loader, [net, distNet], args.cuda)
print('***Test***')
test(test_loader, test_gallery_loader, [net, distNet], args.cuda)
sys.exit()