def run(thread_index, thread_num, result, args):
# initialization
model_dir = os.path.join(args.output_dir, 'model/')
model_path = os.path.join(model_dir, 'latest.pth')
if not osp.exists(model_path):
raise "there is no latest.pth "
output_dir = [model_path, result_dir]
use_cuda = torch.cuda.is_available()
cls_list = ['_' for _ in range(81)]
# datasets
val = TrainDataset(args.base_path,
args.img_list,
'msra',
cls_list,
phase='test')
# val_loader = torch.utils.data.DataLoader(val, batch_size=1, num_workers=1, collate_fn=unique_collate, pin_memory=False)
# model
# model = Encoder_Decoder(args.hidden_size, attn_type=args.attn_type, context_type=args.context_type, n_layers=args.n_layers, multi_head=args.multi_head)
model = Encoder_Decoder(args.hidden_size,
attn_type=args.attn_type,
context_type=args.context_type)
if use_cuda:
model = model.cuda()
model.eval()
thread_result = test_solver(model, val, output_dir, thread_index,
thread_num)
result.extend(thread_result)
# dataset
train = TrainDataset(args.base_path,
args.img_list,
args.use_mode,
cls_list,
weight=args.weight)
train_loader = torch.utils.data.DataLoader(train,
batch_size=1,
num_workers=1,
collate_fn=unique_collate,
pin_memory=False,
shuffle=True)
# model
model = Encoder_Decoder(args.hidden_size,
attn_type=args.attn_type,
context_type=args.context_type)
if args.load:
continue_epochs, continue_iters = load_checkpoint(model, load_pth)
if use_cuda:
model = model.cuda()
if not args.load:
solver(model,
train_loader,
args.num_epochs,
print_every=20,
save_every=10000,
output_dir=output_dir,
learning_rate=args.learning_rate,