def main(): if args.use_cuda: torch.cuda.set_device(args.gpu) dataloader = DataLoader(dict_path=args.dict_path, glove_path=args.glove_path, data_path=args.data_path, batch_size=args.batch_size, use_glove=args.use_glove) model = Net(no_words=dataloader.tokenizer.no_words, lstm_size=args.lstm_size, emb_size=args.emb_size, depth=args.depth) if args.use_cuda: model = model.cuda() if args.start_iter != 0: # load the model state from pre-specified iteration (saved model available) model.load_state_dict(torch.load( os.path.join(args.save_dir, 'iter_%d.pth' % (args.start_iter))), strict=False) tokenizer = Tokenizer(args.dict_path) optimizer = optim.Adam(model.parameters(), lr=args.lr) train(dataloader, model, optimizer, tokenizer)
def load_model_vqa(map_location): global net_vqa global token_to_ix global ix_to_ans # Base path model_path = './model/muan.pkl' token_to_ix_path = './data/token_to_ix.json' pretrained_emb_path = './data/pretrained_emb.npy' ans_dict_path = './data/vqa_answer_dict.json' '''Pre-load''' # Load token_to_ix token_to_ix = json.load(open(token_to_ix_path, 'r')) token_size = len(token_to_ix) print(' ========== Question token vocab size:', token_size) # print('token_to_ix:', token_to_ix) # Load pretrained_emb pretrained_emb = np.load(pretrained_emb_path) print('pretrained_emb shape:', pretrained_emb.shape) # Answers statistic ans_to_ix, ix_to_ans = json.load(open(ans_dict_path, 'r')) ans_size = len(ans_to_ix) print( ' ========== Answer token vocab size (occur more than {} times):'. format(8), ans_size) # print('ix_to_ans:\n', ix_to_ans) '''Load the pre-trained model''' # Load model ckpt time_start = time.time() print('\nLoading ckpt from: {}'.format(model_path)) state_dict = torch.load(model_path, map_location=map_location)['state_dict'] print('state_dict num:', len(state_dict.keys())) print('Finish load state_dict!') # Load model net_vqa = Net(pretrained_emb, token_size, ans_size) net_vqa.load_state_dict(state_dict) net_vqa.cuda() net_vqa.eval() # del state_dict # print('net:', net) time_end = time.time() print('Finish load net model!') print('Model load time: {:.3f}s\n'.format(time_end - time_start))
def run(load_last_checkpoint=False): save_dir = f'{OUTPUT_PATH}/models/' os.makedirs(save_dir, exist_ok=True) neural_net = Net() loss_fn = Loss() optim = torch.optim.SGD(neural_net.parameters(), DEFAULT_LR, momentum=0.9, weight_decay=1e-4) starting_epoch = 0 initial_loss = None if load_last_checkpoint: model_paths = glob(f'''{save_dir}*.ckpt''') model_names = [int(i.split('/')[-1][:-5]) for i in model_paths] latest_model_path = f'''{save_dir}{max(model_names)}.ckpt''' print('loading latest model from:', latest_model_path) checkpoint = torch.load(latest_model_path) neural_net.load_state_dict(checkpoint['model_state_dict']) optim.load_state_dict(checkpoint['optimizer_state_dict']) starting_epoch = checkpoint['epoch'] initial_loss = checkpoint['loss'] if torch.cuda.is_available(): neural_net = neural_net.cuda() loss_fn = loss_fn.cuda() print(f'''Training from epoch: {starting_epoch} towards: {TOTAL_EPOCHS}, with learning rate starting from: {get_lr(starting_epoch)}, and loss: {initial_loss}''') meta = pd.read_csv(f'{OUTPUT_PATH}/augmented_meta.csv', index_col=0).sample(frac=1).reset_index(drop=True) meta_group_by_series = meta.groupby(['seriesuid']).indices list_of_groups = [{i: list(meta_group_by_series[i])} for i in meta_group_by_series.keys()] random.Random(0).shuffle(list_of_groups) val_split = int(VAL_PCT * len(list_of_groups)) val_indices = list(itertools.chain(*[list(i.values())[0] for i in list_of_groups[:val_split]])) train_indices = list(itertools.chain(*[list(i.values())[0] for i in list_of_groups[val_split:]])) ltd = LunaDataSet(train_indices, meta) lvd = LunaDataSet(val_indices, meta) train_loader = DataLoader(ltd, batch_size=1, shuffle=False) val_loader = DataLoader(lvd, batch_size=1, shuffle=False) for ep in range(starting_epoch, TOTAL_EPOCHS): train(train_loader, neural_net, loss_fn, ep, optim, get_lr, save_dir=save_dir) validate(val_loader, neural_net, loss_fn)
print('bboxfeat_batch:', bboxfeat_batch.shape) time_end = time.time() print('Image feature process time: {:.3f}s'.format(time_end - time_start)) # Load model ckpt print('\nLoading ckpt from: {}'.format(model_path)) time_start = time.time() state_dict = torch.load(model_path, map_location=map_location)['state_dict'] print('state_dict num:', len(state_dict.keys())) print('Finish load state_dict!') # Load model net = Net(pretrained_emb, token_size, ans_size) net.cuda() net.eval() net.load_state_dict(state_dict) # print('net:', net) time_end = time.time() print('Finish load net model!') print('Model load time: {:.3f}s\n'.format(time_end - time_start)) # Predict time_start = time.time() pred = net(imgfeat_batch, bboxfeat_batch, quesix_batch) pred_np = pred.cpu().data.numpy() pred_argmax = np.argmax(pred_np, axis=1)[0] pred_ans = ix_to_ans[str(pred_argmax)] print('pred_argmax:', pred_argmax) print('pred_ans:', pred_ans)
args.weight_path, map_location='cuda' if torch.cuda.is_available() else 'cpu') model_dict = model.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v) } model_dict.update(pretrained_dict) model.load_state_dict(model_dict) print('-- Loading weights finished.') # 2.多GPU并行 if torch.cuda.is_available(): model = torch.nn.DataParallel(model) cudnn.benchmark = True model = model.cuda() # 3.创建计算loss的类 criterion = MultiBoxLoss() # 4.创建优化器 optimizer = optim.Adam(model.parameters(), lr=args.lr) model.train() # 5.读取数据开始训练Epoch轮 for epoch in range(args.Epoch): # 5.1每轮使用不同学习率 if epoch % 10 == 0: adjust_learning_rate(optimizer, args.lr, 0.95, epoch) # 5.2创建数据加载器 train_data = MyDataSet(args.annotation_path, Config['input_size'], transform=my_transform, loader=default_loader)
prompt: d.get(prompt)['predictions'] for prompt in d } return self if __name__ == '__main__': args = parser.parse_args() # Loading the evaluation dataset print("Loading dataset") data_params_json_path = os.path.join(args.data_dir, 'params.json') data_params = utils.DataParams.from_json(data_params_json_path) val_dataset = DoulingoDataset(data_params, split='val') # Loading the model print("Loading model...") checkpoint = os.path.join(args.model_dir, f"runs/{args.checkpoint}.pth.tar") config = utils.Params(cuda=torch.cuda.is_available(), src='en', trg='hu') model = Net(config) checkpoint = torch.load(checkpoint) model.load_state_dict(checkpoint['state_dict']) print("Finished Loading") # Evaluation ... print("Starting Evaluation..") if not os.path.exists(args.results_dir): os.mkdir(args.results_dir) metrics = evaluate_model(model.cuda(), val_dataset, args.results_dir) result_json = os.path.join(args.results_dir, 'metrics.json') utils.save_dict_to_json(metrics, result_json)