'_' + str(tta) + '_' + str(ms) + '.pkl')): d = pickle.load(open(pkl_name, 'rb')) pred = d['pred'] else: test_dataloader = data.DataLoader( test_dataset, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, collate_fn=tta_collate if args.use_tta else default_collate, shuffle=False) model = BaseModel(args) model.init_model() model.load_trained_model() pred = model.eval(test_dataloader) # after softmax array pickle.dump({'pred': pred}, open(pkl_name, 'wb')) preds_test += pred preds_test /= len(args.ensemble_exp) # generate csv print('generate csv ...') for t in [0.29, 0.30, 0.31, 0.32, 0.33, 0.4, 0.45]: make_submission((preds_test > t).astype(np.uint8), test['names'], path='{}_{:.2f}_submission.csv'.format(csv_name, t)) print('generate {}_submission.csv'.format(args.exp_name)) if args.vis: print('vis mask ...')
val = pickle.load(open(os.path.join(args.data_root, 'val.pkl'), 'rb')) image_root = os.path.join(args.data_root, 'train', 'images') val_dataset = SaltSet(val, image_root, BaseTransform(args.size, MEAN, STD)) val_dataloader = data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, shuffle=False) model = BaseModel(args) model.init_model() model.load_trained_model() preds_test = model.eval(val_dataloader) pred_dict = { idx: RLenc( np.round( downsample(preds_test[i].squeeze(), img_size_ori, img_size_target) > 0.5)) for i, idx in enumerate(val['names']) } sub = pd.DataFrame.from_dict(pred_dict, orient='index') sub.index.names = ['id'] sub.columns = ['rle_mask'] sub.to_csv('submission.csv')
shuffle=False) pkl_name = os.path.join( ensemble_dict, args.exp_name + '_' + args.trained_model + '_' + str(args.use_tta) + '_' + str(args.ms) + '.pkl') print(pkl_name) if os.path.exists(pkl_name): d = pickle.load(open(pkl_name, 'rb')) preds_test = d['pred'] else: model = BaseModel(args) model.init_model() model.load_trained_model() preds_test = model.eval(test_dataloader) # (bs,101, 101) pickle.dump({'pred': preds_test}, open(pkl_name, 'wb')) # generate csv print('generate csv ...') for t in [ 0.14, ]: make_submission( (preds_test > t).astype(np.uint8), test['names'], path='{}_{:.2f}_submission.csv'.format(args.exp_name, t)) # pred_dict = {idx: RLenc(np.round(preds_test[i] > 0.5)) for # i, idx in tqdm(enumerate(test['names']))} # # sub = pd.DataFrame.from_dict(pred_dict, orient='index')
} # --------------------train and val------------------- best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 print("Start training:") print("Detailed training log path: {}".format("{}log.log".format(work_dir))) for epoch in range(num_epochs): epoch_index = "{}/{}".format(epoch + 1, num_epochs) train_process(dataLoader_train, epoch_index, params) val_acc = val_process(dataLoader_val, epoch_index, params) exp_lr_scheduler.step() if val_acc > best_acc: best_acc = val_acc best_model_wts = copy.deepcopy(model.state_dict()) torch.save(model.state_dict(), "{}epoch-{}.pkl".format(work_dir, epoch)) model.load_state_dict(best_model_wts) torch.save(model.state_dict(), "{}best-val.pkl".format(work_dir)) # --------------------inference------------------- model.load_state_dict(torch.load("{}best-val.pkl".format(work_dir), map_location=lambda storage, loc: storage)) model.eval() print("Start testing:") predict(test_data_path, work_dir, params)