def main(args):
# need to clear the pipeline
# top1 & top10 need to be chosen in the same params ???
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_transform = transforms.Compose([transforms.ToTensor(), normalize])
test_loader = data_config(args.image_dir, args.anno_dir, 64, 'test', 100,
test_transform)
unique_image = get_image_unique(args.image_dir, args.anno_dir, 64, 'test',
100, test_transform)
ac_i2t_top1_best = 0.0
ac_i2t_top10_best = 0.0
ac_t2i_top1_best = 0.0
ac_t2i_top10_best = 0.0
ac_t2i_top5_best = 0.0
ac_i2t_top5_best = 0.0
i2t_models = os.listdir(args.model_path)
i2t_models.sort()
model_list = []
for i2t_model in i2t_models:
if i2t_model.split('.')[0] != "model_best":
model_list.append(int(i2t_model.split('.')[0]))
model_list.sort()
logging.info('Testing on dataset: {}'.format(args.anno_dir))
for i2t_model in model_list:
model_file = os.path.join(args.model_path, str(i2t_model) + '.pth.tar')
if os.path.isdir(model_file):
continue
epoch = i2t_model
if int(epoch) < args.epoch_start:
continue
network, _ = network_config(args, 'test', None, True, model_file)
ac_top1_i2t, ac_top5_i2t, ac_top10_i2t, ac_top1_t2i, ac_top5_t2i, ac_top10_t2i, test_time = test(
test_loader, network, args, unique_image)
if ac_top1_t2i > ac_t2i_top1_best:
ac_i2t_top1_best = ac_top1_i2t
ac_i2t_top5_best = ac_top5_i2t
ac_i2t_top10_best = ac_top10_i2t
ac_t2i_top1_best = ac_top1_t2i
ac_t2i_top5_best = ac_top5_t2i
ac_t2i_top10_best = ac_top10_t2i
dst_best = os.path.join(args.checkpoint_dir, 'model_best',
str(epoch)) + '.pth.tar'
logging.info('epoch:{}'.format(epoch))
logging.info(
'top1_t2i: {:.3f}, top5_t2i: {:.3f}, top10_t2i: {:.3f}, top1_i2t: {:.3f}, top5_i2t: {:.3f}, top10_i2t: {:.3f}'
.format(ac_top1_t2i, ac_top5_t2i, ac_top10_t2i, ac_top1_i2t,
ac_top5_i2t, ac_top10_i2t))
logging.info(
't2i_top1_best: {:.3f}, t2i_top5_best: {:.3f}, t2i_top10_best: {:.3f}, i2t_top1_best: {:.3f}, i2t_top5_best: {:.3f}, i2t_top10_best: {:.3f}'
.format(ac_t2i_top1_best, ac_t2i_top5_best, ac_t2i_top10_best,
ac_i2t_top1_best, ac_i2t_top5_best, ac_i2t_top10_best))
logging.info(args.model_path)
logging.info(args.log_dir)
def main(args):
# transform
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
val_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# data
train_loader = data_config(args.dataset_dir, args.batch_size, 'train',
args.max_length, train_transform)
#val_loader = data_config(args.dataset_dir, 64, 'val', args.max_length, val_transform)
# loss
compute_loss = Loss(args)
nn.DataParallel(compute_loss).cuda()
# network
network, optimizer = network_config(args, 'train',
compute_loss.parameters(), args.resume,
args.model_path)
# lr_scheduler
scheduler = lr_scheduler(optimizer, args)
for epoch in range(args.num_epoches - args.start_epoch):
# train for one epoch
train_loss, train_time, image_precision, text_precision = train(
args.start_epoch + epoch, train_loader, network, optimizer,
compute_loss, args)
# evaluate on validation set
is_best = False
print('Train done for epoch-{}'.format(args.start_epoch + epoch))
state = {
'network': network.state_dict(),
'optimizer': optimizer.state_dict(),
'W': compute_loss.W,
'epoch': args.start_epoch + epoch
}
# 'ac': [ac_top1_i2t, ac_top10_i2t, ac_top1_t2i, ac_top10_t2i],
# 'best_ac': [ac_i2t_best, ac_t2i_best]}
save_checkpoint(state, epoch, args.checkpoint_dir, is_best)
logging.info(
'Epoch: [{}|{}], train_time: {:.3f}, train_loss: {:.3f}'.format(
args.start_epoch + epoch, args.num_epoches, train_time,
train_loss))
logging.info('image_precision: {:.3f}, text_precision: {:.3f}'.format(
image_precision, text_precision))
adjust_lr(optimizer, args.start_epoch + epoch, args)
scheduler.step()
for param in optimizer.param_groups:
print('lr:{}'.format(param['lr']))
break
logging.info('Train done')
logging.info(args.checkpoint_dir)
logging.info(args.log_dir)
def load_model(args):
network, _ = network_config(args, 'test', None, True, model_path, False)
network.eval()
with open(test_sort_path, 'rb') as f_pkl:
test_sort = pickle.load(f_pkl)
with open(word_to_index_path, 'rb') as f_2_pkl:
word_to_index = pickle.load(f_2_pkl)
word_to_index = {k.lower(): v for k, v in word_to_index.items()}
return network, test_sort, word_to_index
def main(args):
# need to clear the pipeline
# top1 & top10 need to be chosen in the same params ???
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
test_loader = data_config(args.image_dir, args.anno_dir, 64, 'test',
args.max_length, test_transform)
ac_i2t_top1_best = 0.0
ac_i2t_top10_best = 0.0
ac_t2i_top1_best = 0.0
ac_t2i_top10_best = 0.0
i2t_models = os.listdir(args.model_path)
i2t_models.sort()
for i2t_model in i2t_models:
model_file = os.path.join(args.model_path, i2t_model)
if os.path.isdir(model_file):
continue
epoch = i2t_model.split('.')[0]
if int(epoch) >= args.epoch_ema:
ema = True
else:
ema = False
network, _ = network_config(args, 'test', None, True, model_file, ema)
ac_top1_i2t, ac_top10_i2t, ac_top1_t2i, ac_top10_t2i, test_time = test(
test_loader, network, args)
if ac_top1_t2i > ac_t2i_top1_best:
ac_i2t_top1_best = ac_top1_i2t
ac_i2t_top10_best = ac_top10_i2t
ac_t2i_top1_best = ac_top1_t2i
ac_t2i_top10_best = ac_top10_t2i
dst_best = os.path.join(args.model_path, 'model_best',
str(epoch)) + '.pth.tar'
shutil.copyfile(model_file, dst_best)
logging.info('epoch:{}'.format(epoch))
logging.info(
'top1_t2i: {:.3f}, top10_t2i: {:.3f}, top1_i2t: {:.3f}, top10_i2t: {:.3f}'
.format(ac_top1_t2i, ac_top10_t2i, ac_top1_i2t, ac_top10_i2t))
logging.info(
't2i_top1_best: {:.3f}, t2i_top10_best: {:.3f}, i2t_top1_best: {:.3f}, i2t_top10_best: {:.3f}'
.format(ac_t2i_top1_best, ac_t2i_top10_best, ac_i2t_top1_best,
ac_i2t_top10_best))
logging.info(args.model_path)
logging.info(args.log_dir)
def main(args):
# transform
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
val_transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
test_transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
cap_transform = None
# data
train_loader = data_config(args.image_dir, args.anno_dir, args.batch_size, 'train', 100, train_transform, cap_transform=cap_transform)
test_loader = data_config(args.image_dir, args.anno_dir, 64, 'test', 100, test_transform)
unique_image = get_image_unique(args.image_dir, args.anno_dir, 64, 'test', 100, test_transform)
# loss
compute_loss = Loss(args)
nn.DataParallel(compute_loss).cuda()
# network
network, optimizer = network_config(args, 'train', compute_loss.parameters(), args.resume, args.model_path)
# lr_scheduler
scheduler = WarmupMultiStepLR(optimizer, (20, 25, 35), 0.1, 0.01, 10, 'linear')
ac_t2i_top1_best = 0.0
best_epoch = 0
for epoch in range(args.num_epoches - args.start_epoch):
network.train()
# train for one epoch
train_loss, train_time, image_precision, text_precision = train(args.start_epoch + epoch, train_loader, network, optimizer, compute_loss, args)
# evaluate on validation set
is_best = False
print('Train done for epoch-{}'.format(args.start_epoch + epoch))
logging.info('Epoch: [{}|{}], train_time: {:.3f}, train_loss: {:.3f}'.format(args.start_epoch + epoch, args.num_epoches, train_time, train_loss))
logging.info('image_precision: {:.3f}, text_precision: {:.3f}'.format(image_precision, text_precision))
scheduler.step()
for param in optimizer.param_groups:
print('lr:{}'.format(param['lr']))
if epoch >= 0:
ac_top1_i2t, ac_top5_i2t, ac_top10_i2t, ac_top1_t2i, ac_top5_t2i , ac_top10_t2i, test_time = test(test_loader, network, args, unique_image)
state = {'network': network.state_dict(), 'optimizer': optimizer.state_dict(), 'W': compute_loss.W, 'epoch': args.start_epoch + epoch}
if ac_top1_t2i > ac_t2i_top1_best:
best_epoch = epoch
ac_t2i_top1_best = ac_top1_t2i
save_checkpoint(state, epoch, args.checkpoint_dir, is_best)
logging.info('epoch:{}'.format(epoch))
logging.info('top1_t2i: {:.3f}, top5_t2i: {:.3f}, top10_t2i: {:.3f}, top1_i2t: {:.3f}, top5_i2t: {:.3f}, top10_i2t: {:.3f}'.format(
ac_top1_t2i, ac_top5_t2i, ac_top10_t2i, ac_top1_i2t, ac_top5_i2t, ac_top10_i2t))
logging.info('Best epoch:{}'.format(best_epoch))
logging.info('Train done')
logging.info(args.checkpoint_dir)
logging.info(args.log_dir)
def load_model():
# global graph
# graph = tf.get_default_graph()
global model
model_file = args.model_path
model, _ = network_config(args, 'test', None, True, model_file, True)
