def inference(): model = DFSeg_model.RedNet(num_classes=40, pretrained=False) #model = nn.DataParallel(model) load_ckpt(model, None, args.last_ckpt, device) model.eval() model.to(device) val_data = SUNRGBD(transform=torchvision.transforms.Compose([scaleNorm(), ToTensor(), Normalize()]), phase_train=False, data_dir=args.data_dir ) val_loader = DataLoader(val_data, batch_size=1, shuffle=False,num_workers=1, pin_memory=True) acc_meter = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() a_meter = AverageMeter() b_meter = AverageMeter() with torch.no_grad(): for batch_idx, sample in enumerate(val_loader): #origin_image = sample['origin_image'].numpy() #origin_depth = sample['origin_depth'].numpy() image = sample['image'].to(device) depth = sample['depth'].to(device) label = sample['label'].numpy() with torch.no_grad(): pred = model(image, depth) output = torch.max(pred, 1)[1] + 1 output = output.squeeze(0).cpu().numpy() acc, pix = accuracy(output, label) intersection, union = intersectionAndUnion(output, label, args.num_class) acc_meter.update(acc, pix) a_m, b_m = macc(output, label, args.num_class) intersection_meter.update(intersection) union_meter.update(union) a_meter.update(a_m) b_meter.update(b_m) print('[{}] iter {}, accuracy: {}' .format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), batch_idx, acc)) # img = image.cpu().numpy() # print('origin iamge: ', type(origin_image)) #if args.visualize: # visualize_result(origin_image, origin_depth, label-1, output-1, batch_idx, args) iou = intersection_meter.sum / (union_meter.sum + 1e-10) for i, _iou in enumerate(iou): print('class [{}], IoU: {}'.format(i, _iou)) mAcc = (a_meter.average() / (b_meter.average()+1e-10)) print(mAcc.mean()) print('[Eval Summary]:') print('Mean IoU: {:.4}, Accuracy: {:.2f}%' .format(iou.mean(), acc_meter.average() * 100))
def my_test(): writer_image = SummaryWriter(os.path.join(args.summary_dir, 'segtest')) model = MultiTaskCNN(38, depth_channel=1, pretrained=False, arch='resnet50', use_aspp=False) load_ckpt(model, None, args.last_ckpt, device) model.eval() model = model.to(device) val_data = my_data_eval.ReadData(transform=torchvision.transforms.Compose( [my_data_eval.scaleNorm(), my_data_eval.ToTensor(), Normalize()]), data_dir=args.data_dir) val_loader = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4, pin_memory=False) with torch.no_grad(): for batch_idx, sample in enumerate(val_loader): # origin_image = sample['origin_image'].numpy() # origin_depth = sample['origin_depth'].numpy() image = sample['image'].to(device) depth = sample['depth'].to(device) with torch.no_grad(): time1 = time.time() pred = model(image, depth) time2 = time.time() print('推理时间:', time2 - time1, '\nFPS: ', 1 / (time2 - time1)) output = torch.max(pred, 1)[1] # output = output.squeeze(0).cpu().numpy() output = output.cpu().numpy() grid_image1 = make_grid(image[:1].clone().cpu().data, 1, normalize=True) writer_image.add_image('image', grid_image1, batch_idx) grid_image2 = make_grid(depth[:1].clone().cpu().data, 1, normalize=True) writer_image.add_image('depth', grid_image2, batch_idx) grid_image3 = make_grid(utils.color_label( torch.max(pred[:1], 1)[1]), 1, normalize=False, range=(0, 255)) writer_image.add_image('Predicted label', grid_image3, batch_idx)
def test(load_model_weight=False): if load_model_weight: if cfg.model_weight_file != '': map_location = (lambda storage, loc: storage) sd = torch.load(cfg.model_weight_file, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format(cfg.model_weight_file)) else: load_ckpt(modules_optims, cfg.ckpt_file) for test_set, name in zip(test_sets, test_set_names): feature_map = ExtractFeature(model_w, TVT) test_set.set_feat_func(feature_map) print('\n=========> Test on dataset: {} <=========\n'.format(name)) test_set.eval( normalize_feat=cfg.normalize_feature, verbose=True)
def inference(): model = ACNet_models_V1.ACNet(num_class=5, pretrained=False) load_ckpt(model, None, None, args.last_ckpt, device) model.eval() model.to(device) data = ACNet_data.FreiburgForest(transform=torchvision.transforms.Compose([ ACNet_data.ScaleNorm(), ACNet_data.ToTensor(), ACNet_data.Normalize() ]), data_dirs=[args.data_dir], modal1_name=args.modal1, modal2_name=args.modal2, gt_available=False) data_loader = DataLoader(data, batch_size=1, shuffle=False, num_workers=1, pin_memory=True) with torch.no_grad(): for batch_idx, sample in enumerate(data_loader): modal1 = sample['modal1'].to(device) modal2 = sample['modal2'].to(device) basename = sample['basename'][0] with torch.no_grad(): pred = model(modal1, modal2) output = torch.argmax(pred, 1) + 1 output = output.squeeze(0).cpu().numpy() if args.save_predictions: colored_output = utils.color_label_eval(output).astype( np.uint8) imageio.imwrite(f'{args.output_dir}/{basename}_pred.png', colored_output.transpose([1, 2, 0]))
def inference(): model = RedNet_model.RedNet(pretrained=False) load_ckpt(model, None, args.last_ckpt, device) model.eval() model.to(device) image = imageio.imread(args.rgb) depth = imageio.imread(args.depth) # Bi-linear image = skimage.transform.resize(image, (image_h, image_w), order=1, mode='reflect', preserve_range=True) # Nearest-neighbor depth = skimage.transform.resize(depth, (image_h, image_w), order=0, mode='reflect', preserve_range=True) image = image / 255 image = torch.from_numpy(image).float() depth = torch.from_numpy(depth).float() image = image.permute(2, 0, 1) depth.unsqueeze_(0) image = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image) depth = torchvision.transforms.Normalize(mean=[19050], std=[9650])(depth) image = image.to(device).unsqueeze_(0) depth = depth.to(device).unsqueeze_(0) pred = model(image, depth) output = utils.color_label(torch.max(pred, 1)[1] + 1)[0] imageio.imsave(args.output, output.cpu().numpy().transpose((1, 2, 0)))
def train_seq2seq(model, dataloaders): train_dataloader, dev_dataloader, test_dataloader = dataloaders criterion = nn.CrossEntropyLoss(ignore_index=constant.pad_idx) if constant.optim == 'Adam': opt = torch.optim.Adam(model.parameters(), lr=constant.lr) elif constant.optim == 'SGD': opt = torch.optim.SGD(model.parameters(), lr=constant.lr) else: print("Optim is not defined") exit(1) start_epoch = 1 if constant.restore: model, opt, start_epoch = load_ckpt(model, opt, constant.restore_path) if constant.USE_CUDA: model.cuda() if constant.embeddings_cpu: model.encoder.embedding.cpu() best_dev = 10000 best_path = '' patience = 3 scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, 'min', factor=0.5, patience=0, min_lr=1e-6) try: for e in range(start_epoch, constant.epochs): model.train() loss_log = [] ppl_log = [] if constant.grid_search: pbar = enumerate(train_dataloader) else: pbar = tqdm(enumerate(train_dataloader), total=len(train_dataloader)) for b, (dialogs, lens, targets, _, _, _, _, _, _) in pbar: if len(train_dataloader) % (b + 1) == 10: torch.cuda.empty_cache() opt.zero_grad() try: batch_size, max_target_len = targets.shape probs = model(dialogs, lens, targets) # Masked CEL trick: Reshape probs to (B*L, V) and targets to (B*L,) and ignore pad idx probs = probs.transpose(0, 1).contiguous().view( batch_size * max_target_len, -1) targets = targets.contiguous().view(batch_size * max_target_len) loss = criterion(probs, targets) # if constant.embeddings_cpu and constant.USE_CUDA: # targets = targets.cuda() # target_lens = target_lens.cuda() # loss = masked_cross_entropy(probs.transpose(0, 1).contiguous(), targets.contiguous(), target_lens) loss.backward() opt.step() ## logging loss_log.append(loss.item()) ppl_log.append(math.exp(loss_log[-1])) if not constant.grid_search: pbar.set_description( "(Epoch {}) TRAIN LOSS:{:.4f} TRAIN PPL:{:.1f}". format(e, np.mean(loss_log), np.mean(ppl_log))) except RuntimeError as err: if 'out of memory' in str(err): print('| WARNING: ran out of memory, skipping batch') torch.cuda.empty_cache() else: raise err ## LOG dev_loss, dev_ppl = eval_seq2seq(model, dev_dataloader, bleu=False) print("(Epoch {}) DEV LOSS: {:.4f}, DEV PPL: {:.1f}".format( e, dev_loss, dev_ppl)) scheduler.step(dev_loss) if (dev_loss < best_dev): best_dev = dev_loss # save best model path = 'trained/data-{}.task-seq2seq.lr-{}.emb-{}.D-{}.H-{}.attn-{}.bi-{}.parse-{}.loss-{}' # lr.embedding.D.H.attn.bi.parse.metric path = path.format(constant.data, constant.lr, constant.embedding, constant.D, constant.H, constant.attn, constant.bi, constant.parse, best_dev) if constant.topk: path += '.topk-{}.tau-{}'.format(constant.topk_size, constant.tau) if constant.grid_search: path += '.grid' best_path = save_model(model, 'loss', best_dev, path) patience = 3 else: patience -= 1 if patience == 0: break if best_dev == 0.0: break except KeyboardInterrupt: if not constant.grid_search: print("KEYBOARD INTERRUPT: Save CKPT and Eval") save = True if input('Save ckpt? (y/n)\t') in [ 'y', 'Y', 'yes', 'Yes' ] else False if save: save_path = save_ckpt(model, opt, e) print("Saved CKPT path: ", save_path) # ask if eval do_eval = True if input('Proceed with eval? (y/n)\t') in [ 'y', 'Y', 'yes', 'Yes' ] else False if do_eval: dev_loss, dev_ppl, dev_bleu, dev_bleus = eval_seq2seq( model, dev_dataloader, bleu=True, beam=constant.beam) print("DEV LOSS: {:.4f}, DEV PPL: {:.1f}, DEV BLEU: {:.4f}". format(dev_loss, dev_ppl, dev_bleu)) print( "BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}" .format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) exit(1) # load and report best model on test torch.cuda.empty_cache() model = load_model(model, best_path) if constant.USE_CUDA: model.cuda() # train_loss, train_ppl, train_bleu, train_bleus = eval_seq2seq(model, train_dataloader, bleu=True, beam=constant.beam) dev_loss, dev_ppl, dev_bleu, dev_bleus = eval_seq2seq(model, dev_dataloader, bleu=True, beam=constant.beam) test_loss, test_ppl, test_bleu, test_bleus = eval_seq2seq( model, test_dataloader, bleu=True, beam=constant.beam) # print("BEST TRAIN LOSS: {:.4f}, TRAIN PPL: {:.1f}, TRAIN BLEU: {:.4f}".format(train_loss, train_ppl, train_bleu)) # print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}".format(train_bleus[0], train_bleus[1], train_bleus[2], train_bleus[3])) print("BEST DEV LOSS: {:.4f}, DEV PPL: {:.1f}, DEV BLEU: {:.4f}".format( dev_loss, dev_ppl, dev_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}". format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print("BEST TEST LOSS: {:.4f}, TEST PPL: {:.1f}, TEST BLEU: {:.4f}".format( test_loss, test_ppl, test_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}". format(test_bleus[0], test_bleus[1], test_bleus[2], test_bleus[3]))
def train_rl(model, dataloaders): train_dataloader, dev_dataloader, test_dataloader = dataloaders clf_criterion = nn.BCEWithLogitsLoss() mle_criterion = nn.CrossEntropyLoss(ignore_index=constant.pad_idx) baseline_criterion = nn.MSELoss() if constant.optim == 'Adam': opt = torch.optim.Adam(model.parameters(), lr=constant.lr) elif constant.optim == 'SGD': opt = torch.optim.SGD(model.parameters(), lr=constant.lr) else: print("Optim is not defined") exit(1) start_epoch = 1 if constant.restore: model, opt, start_epoch = load_ckpt(model, opt, constant.restore_path) if constant.USE_CUDA: model.cuda() best_dev = 0 best_path = '' patience = 3 scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, 'max', factor=0.5, patience=0, min_lr=1e-6) tau = constant.tau tau_min = 0.2 tau_dec = 0.2 pretrain_curiosity = constant.lambda_aux if constant.pretrain_curiosity: pretrain_curiosity = 0.0 try: for e in range(start_epoch, constant.epochs): model.train() reward_log = [] ori_reward_log = [] aux_reward_log = [] # for sentiment agreement / curiosity inv_loss_log = [] # for curiosity f1_log = [] # pretrain curiosity only for first epoch if e > start_epoch: pretrain_curiosity = constant.lambda_aux # temperature annealing if constant.use_tau_anneal and e > start_epoch and constant.tau > tau_min: constant.tau -= tau_dec if constant.grid_search: pbar = enumerate(train_dataloader) else: pbar = tqdm(enumerate(train_dataloader), total=len(train_dataloader)) for b, (dialogs, lens, targets, _, _, sentiments, sentiments_b, _, _) in pbar: if len(train_dataloader) % (b + 1) == 10: torch.cuda.empty_cache() opt.zero_grad() try: B, T = targets.shape if constant.use_self_critical: step_loss, dec_lens_var, R_g, R, greedy_sents, sampled_sents = model( dialogs, lens, targets) # (R_s - R_g) * step_loss rl_loss = torch.mean( torch.sum((R.detach() - R_g.detach()) * step_loss, dim=1) / dec_lens_var.float()) elif constant.use_arl: step_loss, dec_lens_var, rs, R, arl, sampled_sents = model( dialogs, lens, targets) rs = rs.transpose(0, 1).contiguous() rl_loss = (R.detach() - rs.detach()) * step_loss rl_loss = torch.mean( torch.sum(rl_loss * arl, dim=1) / dec_lens_var.float()) else: # probs: (B, T, V), xs: (B, T), R: (B, 1), rs: (B, T) if constant.use_sentiment and constant.aux_reward_model != '': step_loss, dec_lens_var, rs, R_l, R_s, sampled_sents, clf_logits = model( dialogs, lens, targets, sentiments=sentiments) R = constant.lambda_aux * R_l + R_s clf_loss = clf_criterion(clf_logits, sentiments_b) preds = torch.sigmoid(clf_logits.squeeze()) > 0.5 f1 = f1_score(sentiments_b.cpu().numpy(), preds.detach().cpu().numpy(), average='weighted') f1_log.append(f1) elif constant.use_sentiment and constant.use_sentiment_agreement: step_loss, dec_lens_var, rs, R, sampled_sents, clf_logits = model( dialogs, lens, targets, sentiments=sentiments) clf_loss = clf_criterion(clf_logits, sentiments_b) preds = torch.sigmoid(clf_logits.squeeze()) > 0.5 f1 = f1_score(sentiments_b.cpu().numpy(), preds.detach().cpu().numpy(), average='weighted') f1_log.append(f1) elif constant.use_sentiment_agreement: step_loss, dec_lens_var, rs, R, sampled_sents = model( dialogs, lens, targets, sentiments=sentiments) elif constant.use_sentiment: step_loss, dec_lens_var, rs, R, sampled_sents, clf_logits = model( dialogs, lens, targets, sentiments=sentiments) clf_loss = clf_criterion(clf_logits, sentiments_b) preds = torch.sigmoid(clf_logits.squeeze()) > 0.5 f1 = f1_score(sentiments_b.cpu().numpy(), preds.detach().cpu().numpy(), average='weighted') f1_log.append(f1) elif constant.use_curiosity: step_loss, dec_lens_var, rs, R, R_i, L_i, sampled_sents = model( dialogs, lens, targets) rs = rs.transpose(0, 1).contiguous() R_i = R_i.transpose(0, 1).contiguous() baseline_target = R.detach() * R_i.detach() rl_loss = torch.mean( torch.sum( (R.detach() * R_i.detach() - rs.detach()) * step_loss, dim=1) / dec_lens_var.float()) R_i = torch.mean( torch.sum(R_i, dim=1) / dec_lens_var.float()) else: step_loss, dec_lens_var, rs, R, sampled_sents = model( dialogs, lens, targets) if not constant.use_curiosity: # probs = probs.transpose(0, 1).cntiguous() # xs = xs.transpose(0, 1).contiguous() # # (B, T, V) => (B, T) => (B,) # probs = torch.gather(probs, dim=2, index=xs.unsqueeze(2)).squeeze() # probs = -torch.log(probs) rs = rs.transpose(0, 1).contiguous() rl_loss = torch.mean( torch.sum( (R.detach() - rs.detach()) * step_loss, dim=1) / dec_lens_var.float()) if constant.use_hybrid: probs, _ = model(dialogs, lens, targets, use_mle=True) mle_loss = mle_criterion( probs.transpose(0, 1).contiguous().view(B * T, -1), targets.contiguous().view(B * T)) loss = constant.lambda_mle * rl_loss + ( 1 - constant.lambda_mle) * mle_loss elif constant.use_arl: probs, _ = model(dialogs, lens, targets, use_mle=True) arl_c = torch.ones(arl.size()).to(arl.device) - arl mle_criterion.reduction = 'none' mle_loss = mle_criterion( probs.transpose(0, 1).contiguous().view(B * T, -1), targets.contiguous().view(B * T)) mle_loss = torch.mean( torch.sum(mle_loss * arl_c, dim=1)) loss = rl_loss + mle_loss else: loss = rl_loss if constant.use_sentiment: loss = constant.lambda_emo * clf_loss + ( 1 - constant.lambda_emo) * loss if constant.use_curiosity: loss = pretrain_curiosity * loss + ( 1 - constant.beta) * L_i + constant.beta * R_i loss.backward() opt.step() if constant.use_baseline: if constant.use_curiosity: baseline_loss = baseline_criterion( rs, baseline_target) else: # rs (32, T) <==> R (32, 1) baseline_loss = baseline_criterion( rs, tile(R, T, dim=1)) baseline_loss.backward() opt.step() ## logging reward_log.append(torch.mean(R).item()) if constant.use_sentiment and constant.aux_reward_model != '': ori_reward_log.append(torch.mean(R_l).item()) aux_reward_log.append(torch.mean(R_s).item()) if constant.use_curiosity: aux_reward_log.append(torch.mean(R_i).item()) inv_loss_log.append(L_i.item()) if not constant.grid_search: if constant.use_sentiment: if constant.aux_reward_model != '': pbar.set_description( "(Epoch {}) TRAIN R: {:.3f} R_l: {:.3f} R_s: {:.3f} F1: {:.3f}" .format(e, np.mean(reward_log), np.mean(ori_reward_log), np.mean(aux_reward_log), np.mean(f1_log))) else: pbar.set_description( "(Epoch {}) TRAIN REWARD: {:.4f} TRAIN F1: {:.4f}" .format(e, np.mean(reward_log), np.mean(f1_log))) elif constant.use_curiosity: pbar.set_description( "(Epoch {}) TRAIN R: {:.3f} R_i: {:.3f} L_i: {:.3f}" .format(e, np.mean(reward_log), np.mean(aux_reward_log), np.mean(inv_loss_log))) else: pbar.set_description( "(Epoch {}) TRAIN REWARD: {:.4f}".format( e, np.mean(reward_log))) if b % 100 == 0 and b > 0: # if not constant.use_self_critical: # _, greedy_sents = model(dialogs, lens, targets, test=True, use_mle=True) corrects = [ " ".join([ train_dataloader.dataset.lang.index2word[x_t] for x_t in iter(lambda x=iter(gens): next(x), constant.eou_idx) ]) for gens in targets.cpu().data.numpy() ] contexts = [ " ".join([ train_dataloader.dataset.lang.index2word[x_t] for x_t in iter(lambda x=iter(gens): next(x), constant.pad_idx) ]) for gens in dialogs.cpu().data.numpy() ] for d, c, s, r in zip(contexts, corrects, sampled_sents, R.detach().cpu().numpy()): print('reward: ', r) print('dialog: ', d) print('sample: ', s) print('golden: ', c) print() except RuntimeError as err: if 'out of memory' in str(err): print('| WARNING: ran out of memory, skipping batch') torch.cuda.empty_cache() else: print(err) traceback.print_exc() raise err ## LOG if constant.use_sentiment and not constant.use_sentiment_agreement: dev_reward, dev_f1 = eval_rl(model, dev_dataloader, bleu=False) print("(Epoch {}) DEV REWARD: {:.4f}".format(e, dev_reward)) elif constant.use_curiosity: dev_reward, dev_Ri, dev_Li = eval_rl(model, dev_dataloader, bleu=False) print("(Epoch {}) DEV REWARD: {:.3f} R_i: {:.3f} L_i: {:.3f}". format(e, dev_reward, dev_Ri, dev_Li)) else: dev_reward = eval_rl(model, dev_dataloader, bleu=False) print("(Epoch {}) DEV REWARD: {:.4f}".format(e, dev_reward)) scheduler.step(dev_reward) if (dev_reward > best_dev): best_dev = dev_reward # save best model path = 'trained/data-{}.task-rlseq.lr-{}.tau-{}.lambda-{}.reward-{}.{}' path = path.format(constant.data, constant.lr, tau, constant.lambda_mle, best_dev, constant.reward_model.split('/')[1]) if constant.use_curiosity: path += '.curiosity' if constant.aux_reward_model != '': path += '.' + constant.aux_reward_model.split('/')[1] path += '.lambda_aux-{}'.format(constant.lambda_aux) if constant.use_tau_anneal: path += '.tau_anneal' if constant.use_self_critical: path += '.self_critical' if constant.use_current: path += '.current' if constant.use_sentiment: path += '.sentiment' if constant.use_sentiment_agreement: path += '.agreement' if constant.use_context: path += '.context' if constant.topk: path += '.topk-{}'.format(constant.topk_size) if constant.use_arl: path += '.arl' if constant.grid_search: path += '.grid' best_path = save_model(model, 'reward', best_dev, path) patience = 3 else: patience -= 1 if patience == 0: break if constant.aux_reward_model == '' and best_dev == 0.0: break except KeyboardInterrupt: if not constant.grid_search: print("KEYBOARD INTERRUPT: Save CKPT and Eval") save = True if input('Save ckpt? (y/n)\t') in [ 'y', 'Y', 'yes', 'Yes' ] else False if save: save_path = save_ckpt(model, opt, e) print("Saved CKPT path: ", save_path) # ask if eval do_eval = True if input('Proceed with eval? (y/n)\t') in [ 'y', 'Y', 'yes', 'Yes' ] else False if do_eval: if constant.use_sentiment: if constant.aux_reward_model != '': dev_rewards, dev_f1, dev_bleu, dev_bleus = eval_rl( model, dev_dataloader, bleu=True) print( "DEV R: {:.3f} R_l: {:.3f} R_s: {:.3f} DEV F1: {:.3f} DEV B: {:.3f}" .format(dev_rewards[0], dev_rewards[1], dev_rewards[2], dev_f1, dev_bleu)) else: dev_reward, dev_f1, dev_bleu, dev_bleus = eval_rl( model, dev_dataloader, bleu=True) print( "DEV REWARD: {:.4f}, DEV F1: {:.4f}, DEV BLEU: {:.4f}" .format(dev_reward, dev_f1, dev_bleu)) elif constant.use_curiosity: dev_reward, dev_Ri, dev_Li, dev_bleu, dev_bleus = eval_rl( model, dev_dataloader, bleu=True) print( "BEST DEV REWARD: {:.4f} R_i: {:.3f} L_i: {:.3f} BLEU: {:.4f}" .format(dev_reward, dev_Ri, dev_Li, dev_bleu)) else: dev_reward, dev_bleu, dev_bleus = eval_rl(model, dev_dataloader, bleu=True) print("DEV REWARD: {:.4f}, DEV BLEU: {:.4f}".format( dev_reward, dev_bleu)) print( "BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}" .format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) exit(1) # load and report best model on test torch.cuda.empty_cache() model = load_model(model, best_path) if constant.USE_CUDA: model.cuda() if constant.use_sentiment and not constant.use_sentiment_agreement: if constant.aux_reward_model != '': dev_rewards, dev_f1, dev_bleu, dev_bleus = eval_rl(model, dev_dataloader, bleu=True) test_rewards, test_f1, test_bleu, test_bleus = eval_rl( model, test_dataloader, bleu=True) print( "DEV R: {:.3f} R_l: {:.3f} R_s: {:.3f} DEV F1: {:.3f} DEV B: {:.3f}" .format(dev_rewards[0], dev_rewards[1], dev_rewards[2], dev_f1, dev_bleu)) print( "BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}" .format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print( "TEST R: {:.3f} R_l: {:.3f} R_s: {:.3f} TEST F1: {:.3f} TEST B: {:.3f}" .format(test_rewards[0], test_rewards[1], test_rewards[2], test_f1, test_bleu)) print( "BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}" .format(test_bleus[0], test_bleus[1], test_bleus[2], test_bleus[3])) else: dev_reward, dev_f1, dev_bleu, dev_bleus = eval_rl(model, dev_dataloader, bleu=True) test_reward, test_f1, test_bleu, test_bleus = eval_rl( model, test_dataloader, bleu=True) print( "DEV REWARD: {:.4f}, DEV F1: {:.4f}, DEV BLEU: {:.4f}".format( dev_reward, dev_f1, dev_bleu)) print( "BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}" .format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print("TEST REWARD: {:.4f}, TEST F1: {:.4f}, TEST BLEU: {:.4f}". format(test_reward, test_f1, test_bleu)) print( "BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}" .format(test_bleus[0], test_bleus[1], test_bleus[2], test_bleus[3])) elif constant.use_curiosity: dev_reward, dev_Ri, dev_Li, dev_bleu, dev_bleus = eval_rl( model, dev_dataloader, bleu=True) test_reward, test_Ri, test_Li, test_bleu, test_bleus = eval_rl( model, test_dataloader, bleu=True) print("BEST DEV REWARD: {:.4f} R_i: {:.3f} L_i: {:.3f} BLEU: {:.4f}". format(dev_reward, dev_Ri, dev_Li, dev_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}". format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print("BEST TEST REWARD: {:.4f} R_i: {:.3f} L_i: {:.3f} BLEU: {:.4f}". format(test_reward, test_Ri, test_Li, test_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}". format(test_bleus[0], test_bleus[1], test_bleus[2], test_bleus[3])) else: dev_reward, dev_bleu, dev_bleus = eval_rl(model, dev_dataloader, bleu=True) test_reward, test_bleu, test_bleus = eval_rl(model, test_dataloader, bleu=True) print("BEST DEV REWARD: {:.4f}, BLEU: {:.4f}".format( dev_reward, dev_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}". format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print("BEST TEST REWARD: {:.4f}, BLEU: {:.4f}".format( test_reward, test_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}". format(test_bleus[0], test_bleus[1], test_bleus[2], test_bleus[3]))
results.extend(preds.cpu().numpy().tolist()) details.extend(outputs.cpu().detach().numpy().tolist()) print('writing outputs') output_file = Path(args.output) f = output_file.open('w', encoding='utf-8') for filename, result, detail in zip(testset.im_list, results, details): if not isinstance(filename, Path): filename = Path(filename) f.write(f'{filename.name},{scripts[result]}\n') f.close() print(f'finish. writen as {args.output}') if __name__ == "__main__": ckpt = load_ckpt(os.path.join(config['ckpt_dir'], args.ckpt)) if not args.output: args.output = f"{ckpt['params']['model']}_{ckpt['config']['transform']['val']}_{ckpt['config']['input_size']}.txt" transform = get_transforms(ckpt['config']['transform']['val']) print(f'transform: {ckpt["config"]["transform"]["val"]}') print(f'scripts: {get_scripts()}') testset = TestSet(args.test_path, transform) print(f'testset: {args.test_path}') dataloader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=config['num_workers']) num_classes = len(get_scripts())
def train(): train_data = ACNet_data.SUNRGBD(transform=transforms.Compose([ ACNet_data.scaleNorm(), ACNet_data.RandomScale((1.0, 1.4)), ACNet_data.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), ACNet_data.RandomCrop(image_h, image_w), ACNet_data.RandomFlip(), ACNet_data.ToTensor(), ACNet_data.Normalize() ]), phase_train=True, data_dir=args.data_dir) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False) num_train = len(train_data) if args.last_ckpt: model = ACNet_models_V1.ACNet(num_class=40, pretrained=False) else: model = ACNet_models_V1.ACNet(num_class=40, pretrained=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) # CEL_weighted = utils.CrossEntropyLoss2d() CEL_weighted = utils.FocalLoss2d(weight=nyuv2_frq, gamma=2) model.train() model.to(device) CEL_weighted.to(device) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) global_step = 0 if args.last_ckpt: global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device) #hxx for finetuing # lr_decay_lambda = lambda epoch: 0.2 * args.lr_decay_rate ** ((epoch - args.start_epoch) // args.lr_epoch_per_decay) lr_decay_lambda = lambda epoch: (1 - (epoch - args.start_epoch) / (args.epochs - args.start_epoch))**0.9 scheduler = LambdaLR(optimizer, lr_lambda=lr_decay_lambda) writer = SummaryWriter(args.summary_dir) for epoch in range(int(args.start_epoch), args.epochs): # if (epoch - args.start_epoch) % args.lr_epoch_per_decay == 0: scheduler.step(epoch) local_count = 0 last_count = 0 end_time = time.time() if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch: save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch, local_count, num_train) for batch_idx, sample in enumerate(train_loader): image = sample['image'].to(device) depth = sample['depth'].to(device) target_scales = [ sample[s].to(device) for s in ['label', 'label2', 'label3', 'label4', 'label5'] ] optimizer.zero_grad() pred_scales = model(image, depth, args.checkpoint) loss = CEL_weighted(pred_scales, target_scales) loss.backward() optimizer.step() local_count += image.data.shape[0] global_step += 1 if global_step % args.print_freq == 0 or global_step == 1: time_inter = time.time() - end_time count_inter = local_count - last_count print_log(global_step, epoch, local_count, count_inter, num_train, loss, time_inter) end_time = time.time() last_count = local_count save_ckpt(args.ckpt_dir, model, optimizer, global_step, args.epochs, 0, num_train) print("Training completed ")
def train_multitask(model, dataloaders): train_dataloader, dev_dataloader, test_dataloader = dataloaders if constant.optim == 'Adam': opt = torch.optim.Adam(model.parameters(), lr=constant.lr) elif constant.optim == 'SGD': opt = torch.optim.SGD(model.parameters(), lr=constant.lr) else: print("Optim is not defined") exit(1) start_epoch = 1 if constant.restore: model, opt, start_epoch = load_ckpt(model, opt, constant.restore_path) if constant.USE_CUDA: model.cuda() if constant.embeddings_cpu: model.encoder.embedding.cpu() best_gen = 10000 best_emo = 0 best_path = '' patience = 3 scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( opt, 'min', factor=0.5, patience=0, min_lr=1e-6) try: for e in range(start_epoch, constant.epochs): model.train() gen_loss_log = [] emo_loss_log = [] cyc_loss_log = [] ppl_log = [] emo_f1_log = [] cyc_f1_log = [] if constant.grid_search: pbar = enumerate(train_dataloader) else: pbar = tqdm(enumerate(train_dataloader),total=len(train_dataloader)) for b, (dialogs, lens, targets, _, _, emotions, sentiments, _, _) in pbar: if constant.use_sentiment: emotions = sentiments if len(train_dataloader) % (b+1) == 10: torch.cuda.empty_cache() opt.zero_grad() try: # batch_size, max_target_len = targets.shape emo_logits, _ = model(dialogs, lens, targets, emotions=emotions) if emo_logits is not None: emo_loss_log.append(model.loss['emo'].item()) if constant.use_emotion: preds = torch.argmax(emo_logits, dim=1) elif constant.use_sentiment: preds = torch.sigmoid(emo_logits.squeeze()) > 0.5 emo_f1 = f1_score(emotions.cpu().numpy(), preds.detach().cpu().numpy(), average='weighted') emo_f1_log.append(emo_f1) else: emo_loss_log = 100 emo_f1_log = 0 model.backward() opt.step() ## logging gen_loss_log.append(model.loss['gen'].item()) ppl_log.append(math.exp(gen_loss_log[-1])) if not constant.grid_search: pbar.set_description("(Epoch {}) L_G:{:.4f} PPL:{:.1f} L_E:{:.4f} F1_E:{:.4f}".format( e, np.mean(gen_loss_log), np.mean(ppl_log), np.mean(emo_loss_log), np.mean(emo_f1_log))) except RuntimeError as err: if 'out of memory' in str(err): print('| WARNING: ran out of memory, skipping batch') torch.cuda.empty_cache() else: raise err ## LOG (gen_loss, ppl), (emo_f1) = eval_multitask(model, dev_dataloader, bleu=False) print("(Epoch {}) DEV GEN LOSS:{:.4f} DEV PPL:{:.1f} DEV EMO F1:{:.4f}".format(e, gen_loss, ppl, emo_f1)) scheduler.step(gen_loss) if gen_loss < best_gen: best_gen = gen_loss best_emo = emo_f1 # save best model path = 'trained/data-{}.task-multiseq.lr-{}.emb-{}.D-{}.H-{}.attn-{}.bi-{}.parse-{}.gen_loss-{}.emo_f1-{}' # lr.embedding.D.H.attn.bi.parse.metric path = path.format(constant.data, constant.lr, constant.embedding, constant.D, constant.H, constant.attn, constant.bi, constant.parse, best_gen, best_emo) if constant.use_sentiment: path += '.sentiment' if constant.grid_search: path += '.grid' best_path = save_model(model, 'loss', best_gen, path) patience = 3 else: patience -= 1 if patience == 0: break if best_gen == 0.0: break except KeyboardInterrupt: if not constant.grid_search: print("KEYBOARD INTERRUPT: Save CKPT and Eval") save = True if input('Save ckpt? (y/n)\t') in ['y', 'Y', 'yes', 'Yes'] else False if save: save_path = save_ckpt(model, opt, e) print("Saved CKPT path: ", save_path) # ask if eval do_eval = True if input('Proceed with eval? (y/n)\t') in ['y', 'Y', 'yes', 'Yes'] else False if do_eval: (dev_loss, dev_ppl, dev_bleu, dev_bleus), (emo_f1) = eval_multitask(model, dev_dataloader, bleu=True, beam=constant.beam) print("DEV LOSS: {:.4f}, DEV PPL: {:.1f}, DEV BLEU: {:.4f}".format(dev_loss, dev_ppl, dev_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}".format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print("DEV EMO F1: {:.4f}".format(emo_f1)) exit(1) # load and report best model on test torch.cuda.empty_cache() model = load_model(model, best_path) if constant.USE_CUDA: model.cuda() (dev_loss, dev_ppl, dev_bleu, dev_bleus), (dev_emo_f1) = eval_multitask(model, dev_dataloader, bleu=True, beam=constant.beam) (test_loss, test_ppl, test_bleu, test_bleus), (test_emo_f1) = eval_multitask(model, test_dataloader, bleu=True, beam=constant.beam) print("BEST DEV LOSS: {:.4f}, DEV PPL: {:.1f}, DEV BLEU: {:.4f}".format(dev_loss, dev_ppl, dev_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}".format(dev_bleus[0], dev_bleus[1], dev_bleus[2], dev_bleus[3])) print("DEV EMO F1: {:.4f}".format(dev_emo_f1)) print("BEST TEST LOSS: {:.4f}, TEST PPL: {:.1f}, TEST BLEU: {:.4f}".format(test_loss, test_ppl, test_bleu)) print("BLEU 1: {:.4f}, BLEU 2: {:.4f}, BLEU 3: {:.4f}, BLEU 4: {:.4f}".format(test_bleus[0], test_bleus[1], test_bleus[2], test_bleus[3])) print("TEST EMO F1: {:.4f}".format(test_emo_f1))
def train(cfg): ''' This is the main loop for training Loads the dataset, model, and other things ''' print json.dumps(cfg, sort_keys=True, indent=4) use_cuda = cfg['use-cuda'] _, _, train_dl, val_dl = utils.get_data_loaders(cfg) model = utils.get_model(cfg) if use_cuda: model = model.cuda() model = utils.init_weights(model, cfg) # Get pretrained models, optimizers and loss functions optim = utils.get_optimizers(model, cfg) model, optim, metadata = utils.load_ckpt(model, optim, cfg) loss_fn = utils.get_losses(cfg) # Set up random seeds seed = np.random.randint(2**32) ckpt = 0 if metadata is not None: seed = metadata['seed'] ckpt = metadata['ckpt'] # Get schedulers after getting checkpoints scheduler = utils.get_schedulers(optim, cfg, ckpt) # Print optimizer state print optim # Get loss file handle to dump logs to if not os.path.exists(cfg['save-path']): os.makedirs(cfg['save-path']) lossesfile = open(os.path.join(cfg['save-path'], 'losses.txt'), 'a+') # Random seed according to what the saved model is np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) # Run training loop num_epochs = cfg['train']['num-epochs'] for epoch in range(num_epochs): # Run the main training loop model.train() for data in train_dl: # zero out the grads optim.zero_grad() # Change to required device for key, value in data.items(): data[key] = Variable(value) if use_cuda: data[key] = data[key].cuda() # Get all outputs outputs = model(data) loss_val = loss_fn(outputs, data, cfg) # print it print('Epoch: {}, step: {}, loss: {}'.format( epoch, ckpt, loss_val.data.cpu().numpy())) # Log into the file after some epochs if ckpt % cfg['train']['step-log'] == 0: lossesfile.write('Epoch: {}, step: {}, loss: {}\n'.format( epoch, ckpt, loss_val.data.cpu().numpy())) # Backward loss_val.backward() optim.step() # Update schedulers scheduler.step() # Peek into the validation set ckpt += 1 if ckpt % cfg['peek-validation'] == 0: model.eval() with torch.no_grad(): for val_data in val_dl: # Change to required device for key, value in val_data.items(): val_data[key] = Variable(value) if use_cuda: val_data[key] = val_data[key].cuda() # Get all outputs outputs = model(val_data) loss_val = loss_fn(outputs, val_data, cfg) print 'Validation loss: {}'.format( loss_val.data.cpu().numpy()) lossesfile.write('Validation loss: {}\n'.format(\ loss_val.data.cpu().numpy())) utils.save_images(val_data, outputs, cfg, ckpt) break model.train() # Save checkpoint utils.save_ckpt((model, optim), cfg, ckpt, seed) lossesfile.close()
def validate(cfg): ''' Main loop for validation, load the dataset, model, and other things. Run validation on the validation set ''' print json.dumps(cfg, sort_keys=True, indent=4) use_cuda = cfg['use-cuda'] _, _, _, val_dl = utils.get_data_loaders(cfg) model = utils.get_model(cfg) if use_cuda: model = model.cuda() model = utils.init_weights(model, cfg) model.eval() # Get pretrained models, optimizers and loss functions optim = utils.get_optimizers(model, cfg) model, _, metadata = utils.load_ckpt(model, optim, cfg) loss_fn = utils.get_losses(cfg) # Set up random seeds if metadata is not None: seed = metadata['seed'] # Validation code, reproducibility is required seed = 42 # Random seed according to what the saved model is np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) # Run test loop losses_list = [] # Run the main training loop for idx, data in enumerate(val_dl): # Change to required device for key, value in data.items(): data[key] = Variable(value) if use_cuda: data[key] = data[key].cuda() data = utils.repeat_data(data, cfg) # Get all outputs outputs = model(data) loss_val = loss_fn(outputs, data, cfg, val=True) losses_list.append(float(loss_val)) # print it print('Step: {}, val_loss: {}'.format(idx, loss_val.data.cpu().numpy())) if cfg['val']['save-img']: print outputs['out'].shape utils.save_val_images(data, outputs, cfg, idx) print( """ Summary: Mean: {}, Std: {}, 25per: {}, 50per: {}, 75per: {}, """.format( np.mean(losses_list), np.std(losses_list), np.percentile(losses_list, 25), np.percentile(losses_list, 50), np.percentile(losses_list, 75), ))
def train(): # 记录数据在tensorboard中显示 writer = SummaryWriter(args.summary_dir) # 准备数据集 train_data = data_eval.ReadData(transform=transforms.Compose([ data_eval.scaleNorm(), data_eval.RandomScale((1.0, 1.4)), data_eval.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), data_eval.RandomCrop(image_h, image_w), data_eval.RandomFlip(), data_eval.ToTensor(), data_eval.Normalize() ]), data_dir=args.data_dir) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False, drop_last=True) num_train = len(train_data) # data = iter(train_loader).next() # print('data:', data['image'].shape) # build model if args.last_ckpt: model = MultiTaskCNN(38, depth_channel=1, pretrained=False, arch='resnet18') else: model = MultiTaskCNN(38, depth_channel=1, pretrained=True, arch='resnet18') model = model.to(device) # build optimizer if args.optimizer == 'rmsprop': optimizer = torch.optim.RMSprop(model.parameters(), args.lr) elif args.optimizer == 'sgd': optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=0.9, weight_decay=1e-4) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(model.parameters(), args.lr) else: # rmsprop print('not supported optimizer \n') return None global_step = 0 # 如果有模型的训练权重,则获取global_step,start_epoch if args.last_ckpt: global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device) # # 监测使用哪几个GPU # if torch.cuda.device_count() > 1: # print("Let's use", torch.cuda.device_count(), "GPUs!") # # nn.DataParallel(module, device_ids=None, output_device=None, dim=0):使用多块GPU进行计算 # model = nn.DataParallel(model) model.train() # cal_param(model, data) loss_func = nn.CrossEntropyLoss(weight=weight.float()) for epoch in range(int(args.start_epoch), args.epochs): tq = tqdm(total=len(train_loader) * args.batch_size) lr = poly_lr_scheduler(optimizer, args.lr, iter=epoch, max_iter=args.epochs) tq.set_description('epoch %d, lr %f' % (epoch, lr)) loss_record = [] local_count = 0 # print('1') for batch_idx, data in enumerate(train_loader): # print(batch_idx) image = data['image'].to(device) depth = data['depth'].to(device) label = data['label'].long().to(device) # print('label', label.shape) output, output_sup1, output_sup2 = model(image, depth) loss1 = loss_func(output, label) loss2 = loss_func(output_sup1, label) loss3 = loss_func(output_sup2, label) loss = loss1 + loss2 + loss3 tq.update(args.batch_size) tq.set_postfix(loss='%.6f' % loss) optimizer.zero_grad() loss.backward() optimizer.step() global_step += 1 local_count += image.data.shape[0] writer.add_scalar('loss_step', loss, global_step) loss_record.append(loss.item()) if global_step % args.print_freq == 0 or global_step == 1: for name, param in model.named_parameters(): writer.add_histogram(name, param.clone().cpu().data.numpy(), global_step, bins='doane') writer.add_graph(model, [image, depth]) grid_image1 = make_grid(image[:3].clone().cpu().data, 3, normalize=True) writer.add_image('image', grid_image1, global_step) grid_image2 = make_grid(depth[:3].clone().cpu().data, 3, normalize=True) writer.add_image('depth', grid_image2, global_step) grid_image3 = make_grid(utils.color_label( torch.max(output[:3], 1)[1]), 3, normalize=False, range=(0, 255)) writer.add_image('Predicted label', grid_image3, global_step) grid_image4 = make_grid(utils.color_label(label[:3]), 3, normalize=False, range=(0, 255)) writer.add_image('Groundtruth label', grid_image4, global_step) tq.close() loss_train_mean = np.mean(loss_record) writer.add_scalar('epoch/loss_epoch_train', float(loss_train_mean), epoch) print('loss for train : %f' % loss_train_mean) # 每隔save_epoch_freq个epoch就保存一次权重 if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch: if not os.path.isdir(args.ckpt_dir): os.mkdir(args.ckpt_dir) save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch, local_count, num_train)
def train(): train_dirs = [train_dir for train_dir in [args.train_dir, args.train_dir2] if train_dir is not None] train_data = ACNet_data.FreiburgForest( transform=transforms.Compose([ ACNet_data.ScaleNorm(), # ACNet_data.RandomRotate((-13, 13)), # ACNet_data.RandomSkew((-0.05, 0.10)), ACNet_data.RandomScale((1.0, 1.4)), ACNet_data.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), ACNet_data.RandomCrop(image_h, image_w), ACNet_data.RandomFlip(), ACNet_data.ToTensor(), ACNet_data.Normalize() ]), data_dirs=train_dirs, modal1_name=args.modal1, modal2_name=args.modal2, ) valid_dirs = [valid_dir for valid_dir in [args.valid_dir, args.valid_dir2] if valid_dir is not None] valid_data = ACNet_data.FreiburgForest( transform=transforms.Compose([ ACNet_data.ScaleNorm(), ACNet_data.ToTensor(), ACNet_data.Normalize() ]), data_dirs=valid_dirs, modal1_name=args.modal1, modal2_name=args.modal2, ) ''' # Split dataset into training and validation dataset_length = len(data) valid_split = 0.05 # tiny split due to the small size of the dataset valid_length = int(valid_split * dataset_length) train_length = dataset_length - valid_length train_data, valid_data = torch.utils.data.random_split(data, [train_length, valid_length]) ''' train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False) valid_loader = DataLoader(valid_data, batch_size=args.batch_size * 3, shuffle=False, num_workers=1, pin_memory=False) # Initialize model if args.last_ckpt: model = ACNet_models_V1.ACNet(num_class=5, pretrained=False) else: model = ACNet_models_V1.ACNet(num_class=5, pretrained=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) model.train() model.to(device) # Initialize criterion, optimizer and scheduler criterion = utils.CrossEntropyLoss2d(weight=freiburgforest_frq) criterion.to(device) # TODO: try with different optimizers and schedulers (CyclicLR exp_range for example) # TODO: try with a smaller LR (currently loss decay is too steep and then doesn't change) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # lr_decay_lambda = lambda epoch: args.lr_decay_rate ** (epoch // args.lr_epoch_per_decay) # scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_decay_lambda) scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=int(np.ceil(args.epochs / 7)), T_mult=2, eta_min=8e-4) # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=5e-4) global_step = 0 # TODO: add early stop to avoid overfitting # Continue training from previous checkpoint if args.last_ckpt: global_step, args.start_epoch = utils.load_ckpt(model, optimizer, scheduler, args.last_ckpt, device) writer = SummaryWriter(args.summary_dir) losses = [] for epoch in tqdm(range(int(args.start_epoch), args.epochs)): if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch: utils.save_ckpt(args.ckpt_dir, model, optimizer, scheduler, global_step, epoch) for batch_idx, sample in enumerate(train_loader): modal1, modal2 = sample['modal1'].to(device), sample['modal2'].to(device) target_scales = [sample[s].to(device) for s in ['label', 'label2', 'label3', 'label4', 'label5']] optimizer.zero_grad() pred_scales = model(modal1, modal2, args.checkpoint) loss = criterion(pred_scales, target_scales) loss.backward() optimizer.step() losses.append(loss.item()) global_step += 1 if global_step % args.print_freq == 0 or global_step == 1: for name, param in model.named_parameters(): writer.add_histogram(name, param.detach().cpu().numpy(), global_step, bins='doane') grid_image = make_grid(modal1[:3].detach().cpu(), 3, normalize=False) writer.add_image('Modal1', grid_image, global_step) grid_image = make_grid(modal2[:3].detach().cpu(), 3, normalize=False) writer.add_image('Modal2', grid_image, global_step) grid_image = make_grid(utils.color_label(torch.argmax(pred_scales[0][:3], 1) + 1), 3, normalize=True, range=(0, 255)) writer.add_image('Prediction', grid_image, global_step) grid_image = make_grid(utils.color_label(target_scales[0][:3]), 3, normalize=True, range=(0, 255)) writer.add_image('GroundTruth', grid_image, global_step) writer.add_scalar('Loss', loss.item(), global_step=global_step) writer.add_scalar('Loss average', sum(losses) / len(losses), global_step=global_step) writer.add_scalar('Learning rate', scheduler.get_last_lr()[0], global_step=global_step) # Compute validation metrics with torch.no_grad(): model.eval() losses_val = [] acc_list = [] iou_list = [] for sample_val in valid_loader: modal1_val, modal2_val = sample_val['modal1'].to(device), sample_val['modal2'].to(device) target_val = sample_val['label'].to(device) pred_val = model(modal1_val, modal2_val) losses_val.append(criterion([pred_val], [target_val]).item()) acc_list.append(utils.accuracy( (torch.argmax(pred_val, 1) + 1).detach().cpu().numpy().astype(int), target_val.detach().cpu().numpy().astype(int))[0]) iou_list.append(utils.compute_IoU( y_pred=(torch.argmax(pred_val, 1) + 1).detach().cpu().numpy().astype(int), y_true=target_val.detach().cpu().numpy().astype(int), num_classes=5 )) writer.add_scalar('Loss validation', sum(losses_val) / len(losses_val), global_step=global_step) writer.add_scalar('Accuracy', sum(acc_list) / len(acc_list), global_step=global_step) iou = np.mean(np.stack(iou_list, axis=0), axis=0) writer.add_scalar('IoU_Road', iou[0], global_step=global_step) writer.add_scalar('IoU_Grass', iou[1], global_step=global_step) writer.add_scalar('IoU_Vegetation', iou[2], global_step=global_step) writer.add_scalar('IoU_Sky', iou[3], global_step=global_step) writer.add_scalar('IoU_Obstacle', iou[4], global_step=global_step) writer.add_scalar('mIoU', np.mean(iou), global_step=global_step) model.train() losses = [] scheduler.step() utils.save_ckpt(args.ckpt_dir, model, optimizer, scheduler, global_step, args.epochs) print("Training completed ")
model = RegressionModel(**cfg.model_kwargs) # Loss function criterion = nn.MSELoss() # Optimizer optimizer = optim.Adamax(model.parameters(), lr=cfg.lr) # Learning rate scheduler; change the learning rate dynamically scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer = optimizer, milestones = cfg.epochs_per_decay, \ gamma=cfg.staircase_decay_multiple_factor) # bind the model and optimizer; can be moved to cpu or gpu togethre modules_optims = [model, optimizer] # Load existing model weight or not if cfg.load_model_weight: # load the model weights, optimizer, start epoch from the saved model files start_epoch, scores = load_ckpt(modules_optims, cfg.ckpt_file) else: start_epoch = 0 # transfer model to cpu or gpu; -1 for cpu; >=0 for gpu transfer_optim_state(state=optimizer.state, device_id=-1) LowestLoss = math.inf # the lowest mean squared error HighestR2 = 0 # the highest r^2 value best_epoch = 0 # the index of the epoch achieving the best performance # iterate over all epochs for epoch in range(start_epoch, cfg.total_epochs): # set the model to "train" for training may_set_mode(modules_optims, 'train') # recording loss, time
def inference(): useGpu = True fileName = "run_" if not useGpu: fileName += "cpu_" device = torch.device( "cuda:0" if useGpu and torch.cuda.is_available() else "cpu") os.makedirs('results', exist_ok=True) f = open( "results/" + fileName + str(int(round(time.time() * 1000))) + ".txt", "w+") f.write('=== Start time: ' + str(datetime.now()) + '\n') p = psutil.Process(os.getpid()) nvidia_smi.nvmlInit() handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0) model = RedNet_model.RedNet(pretrained=False) load_ckpt(model, None, args.last_ckpt, device) model.eval() model.to(device) print('Starting list image files') filesCount = 0 files = glob.glob("datasets/mestrado/**/rgb/*.png", recursive=True) files.extend(glob.glob("datasets/mestrado/**/rgb/*.jpg", recursive=True)) cpuTimes = [0.0, 0.0, 0.0, 0.0] gpuTimes = 0.0 gpuMemTimes = 0.0 maxNumThreads = 0 memUsageTimes = 0 for imagePath in files: print('imagePath: ' + imagePath) pathRgb = Path(imagePath) datasetName = osp.basename(str(pathRgb.parent.parent)) # print('datasetName: ' + datasetName) parentDatasetDir = str(pathRgb.parent.parent) # print('parentDatasetDir: ' + parentDatasetDir) depthImageName = os.path.basename(imagePath).replace('jpg', 'png') image = imageio.imread(imagePath) depth = imageio.imread(parentDatasetDir + '/depth/' + depthImageName) if datasetName == "active_vision" or datasetName == "putkk": image = image[0:1080, 240:1680] depth = depth[0:1080, 240:1680] elif datasetName == "semantics3d_mod": image = image[270:1080, 0:1080] depth = depth[270:1080, 0:1080] elif datasetName == "semantics3d_raw": image = image[64:1024, 0:1280] depth = depth[64:1024, 0:1280] # Bi-linear image = skimage.transform.resize(image, (image_h, image_w), order=1, mode='reflect', preserve_range=True) # Nearest-neighbor depth = skimage.transform.resize(depth, (image_h, image_w), order=0, mode='reflect', preserve_range=True) image = image / 255 image = torch.from_numpy(image).float() depth = torch.from_numpy(depth).float() image = image.permute(2, 0, 1) depth.unsqueeze_(0) image = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image) depth = torchvision.transforms.Normalize(mean=[19050], std=[9650])(depth) image = image.to(device).unsqueeze_(0) depth = depth.to(device).unsqueeze_(0) pred = model(image, depth) res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle) mem_res = nvidia_smi.nvmlDeviceGetMemoryInfo(handle) curGpuTime = res.gpu #curGpuMemTime = res.memory #(in percent) curGpuMemTime = mem_res.used / 1e+6 gpuTimes += curGpuTime gpuMemTimes += curGpuMemTime f.write('GPU Usage Percent: ' + str(curGpuTime) + '\n') f.write('GPU Mem Usage (MB)): ' + str(curGpuMemTime) + '\n') curProcessCpuPerU = p.cpu_percent() curCpusPerU = psutil.cpu_percent(interval=None, percpu=True) # gives a single float value for i in range(len(cpuTimes)): curProcessCpu = curProcessCpuPerU curCpu = curCpusPerU[i] cpuTimes[i] += curCpu f.write('Process CPU Percent: ' + str(curProcessCpu) + ' --- CPU Percent: ' + str(curCpu) + '\n') # you can convert that object to a dictionary memInfo = dict(p.memory_full_info()._asdict()) curMemUsage = memInfo['uss'] memUsageTimes += curMemUsage f.write('Process memory usage: ' + str(curMemUsage / 1e+6) + '\n') f.write('Memory information: ' + str(memInfo) + '\n') if maxNumThreads < p.num_threads(): maxNumThreads = p.num_threads() # print('############## Index: ') # print(index) os.makedirs('results/' + datasetName, exist_ok=True) output = utils.to_label(torch.max(pred, 1)[1] + 1) #output = utils.to_label(torch.max(pred, 1)[1] + 1)[0] #imageio.imsave('results/' + datasetName + '/' + depthImageName, output.cpu().numpy().transpose((1, 2, 0))) #imageio.imsave('results/' + datasetName + '/' + depthImageName, output) lbl_pil = PIL.Image.fromarray(output.astype(np.uint8), mode='P') lbl_pil.save('results/' + datasetName + '/' + depthImageName) filesCount = filesCount + 1 del image, depth, pred, output torch.cuda.empty_cache() nvidia_smi.nvmlShutdown() start = time.time() for imagePath in files: pathRgb = Path(imagePath) datasetName = osp.basename(str(pathRgb.parent.parent)) parentDatasetDir = str(pathRgb.parent.parent) depthImageName = os.path.basename(imagePath).replace('jpg', 'png') image = imageio.imread(imagePath) depth = imageio.imread(parentDatasetDir + '/depth/' + depthImageName) if datasetName == "active_vision" or datasetName == "putkk": image = image[0:1080, 240:1680] depth = depth[0:1080, 240:1680] elif datasetName == "semantics3d_mod": image = image[270:1080, 0:1080] depth = depth[270:1080, 0:1080] elif datasetName == "semantics3d_raw": image = image[64:1024, 0:1280] depth = depth[64:1024, 0:1280] # Bi-linear image = skimage.transform.resize(image, (image_h, image_w), order=1, mode='reflect', preserve_range=True) # Nearest-neighbor depth = skimage.transform.resize(depth, (image_h, image_w), order=0, mode='reflect', preserve_range=True) image = image / 255 image = torch.from_numpy(image).float() depth = torch.from_numpy(depth).float() image = image.permute(2, 0, 1) depth.unsqueeze_(0) image = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image) depth = torchvision.transforms.Normalize(mean=[19050], std=[9650])(depth) image = image.to(device).unsqueeze_(0) depth = depth.to(device).unsqueeze_(0) pred = model(image, depth) del image, depth, pred #torch.cuda.empty_cache() end = time.time() f.write('=== Mean GPU Usage Percent: ' + str(gpuTimes / filesCount) + '\n') f.write('=== Mean GPU Mem Usage (MB): ' + str(gpuMemTimes / filesCount) + '\n') for i in range(len(cpuTimes)): f.write("=== Mean cpu" + str(i) + " usage: " + str(cpuTimes[i] / filesCount) + '\n') f.write("=== Mean memory usage (MB): " + str((memUsageTimes / filesCount) / 1e+6) + '\n') f.write("=== Total image predicted: " + str(filesCount) + '\n') f.write("=== Seconds per image: " + str(((end - start) / filesCount)) + '\n') f.write("=== Max num threads: " + str(maxNumThreads) + '\n') f.write('=== End time: ' + str(datetime.now()) + '\n') f.close()
def evaluate(): model = ACNet_models_V1.ACNet(num_class=5, pretrained=False) load_ckpt(model, None, None, args.last_ckpt, device) model.eval() model.to(device) val_data = ACNet_data.FreiburgForest( transform=torchvision.transforms.Compose([ ACNet_data.ScaleNorm(), ACNet_data.ToTensor(), ACNet_data.Normalize() ]), data_dirs=[args.test_dir], modal1_name=args.modal1, modal2_name=args.modal2, ) val_loader = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=1, pin_memory=True) acc_meter = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() a_meter = AverageMeter() b_meter = AverageMeter() with torch.no_grad(): for batch_idx, sample in enumerate(val_loader): modal1 = sample['modal1'].to(device) modal2 = sample['modal2'].to(device) label = sample['label'].numpy() basename = sample['basename'][0] with torch.no_grad(): pred = model(modal1, modal2) output = torch.argmax(pred, 1) + 1 output = output.squeeze(0).cpu().numpy() acc, pix = accuracy(output, label) intersection, union = intersectionAndUnion(output, label, args.num_class) acc_meter.update(acc, pix) a_m, b_m = macc(output, label, args.num_class) intersection_meter.update(intersection) union_meter.update(union) a_meter.update(a_m) b_meter.update(b_m) print('[{}] iter {}, accuracy: {}' .format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), batch_idx, acc)) if args.visualize: visualize_result(modal1, modal2, label, output, batch_idx, args) if args.save_predictions: colored_output = utils.color_label_eval(output).astype(np.uint8) imageio.imwrite(f'{args.output_dir}/{basename}_pred.png', colored_output.transpose([1, 2, 0])) iou = intersection_meter.sum / (union_meter.sum + 1e-10) for i, _iou in enumerate(iou): print('class [{}], IoU: {}'.format(i, _iou)) mAcc = (a_meter.average() / (b_meter.average() + 1e-10)) print(mAcc.mean()) print('[Eval Summary]:') print('Mean IoU: {:.4}, Accuracy: {:.2f}%' .format(iou.mean(), acc_meter.average() * 100))
def train(): # 记录数据在tensorboard中显示 writer_loss = SummaryWriter(os.path.join(args.summary_dir, 'loss')) # writer_loss1 = SummaryWriter(os.path.join(args.summary_dir, 'loss', 'loss1')) # writer_loss2 = SummaryWriter(os.path.join(args.summary_dir, 'loss', 'loss2')) # writer_loss3 = SummaryWriter(os.path.join(args.summary_dir, 'loss', 'loss3')) writer_acc = SummaryWriter(os.path.join(args.summary_dir, 'macc')) # 准备数据集 train_data = data_eval.ReadData(transform=transforms.Compose([ data_eval.scaleNorm(), data_eval.RandomScale((1.0, 1.4)), data_eval.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), data_eval.RandomCrop(image_h, image_w), data_eval.RandomFlip(), data_eval.ToTensor(), data_eval.Normalize() ]), data_dir=args.train_data_dir) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False, drop_last=True) val_data = data_eval.ReadData(transform=transforms.Compose([ data_eval.scaleNorm(), data_eval.RandomScale((1.0, 1.4)), data_eval.RandomCrop(image_h, image_w), data_eval.ToTensor(), data_eval.Normalize() ]), data_dir=args.val_data_dir) val_loader = DataLoader(val_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False, drop_last=True) num_train = len(train_data) # num_val = len(val_data) # build model if args.last_ckpt: model = MultiTaskCNN_Atten(38, depth_channel=1, pretrained=False, arch='resnet50', use_aspp=True) else: model = MultiTaskCNN_Atten(38, depth_channel=1, pretrained=True, arch='resnet50', use_aspp=True) # build optimizer if args.optimizer == 'rmsprop': optimizer = torch.optim.RMSprop(model.parameters(), args.lr) elif args.optimizer == 'sgd': optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=0.9, weight_decay=1e-4) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(model.parameters(), args.lr) else: # rmsprop print('not supported optimizer \n') return None global_step = 0 max_miou_val = 0 loss_count = 0 # 如果有模型的训练权重,则获取global_step,start_epoch if args.last_ckpt: global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device) # if torch.cuda.device_count() > 1 and args.cuda and torch.cuda.is_available(): # print("Let's use", torch.cuda.device_count(), "GPUs!") # model = torch.nn.DataParallel(model).to(device) model = model.to(device) model.train() # cal_param(model, data) loss_func = nn.CrossEntropyLoss() for epoch in range(int(args.start_epoch), args.epochs): torch.cuda.empty_cache() # if epoch <= freeze_epoch: # for layer in [model.conv1, model.maxpool,model.layer1, model.layer2, model.layer3, model.layer4]: # for param in layer.parameters(): # param.requires_grad = False tq = tqdm(total=len(train_loader) * args.batch_size) if loss_count >= 10: args.lr = 0.5 * args.lr loss_count = 0 lr = poly_lr_scheduler(optimizer, args.lr, iter=epoch, max_iter=args.epochs) optimizer.param_groups[0]['lr'] = lr # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 30, gamma=0.5) tq.set_description('epoch %d, lr %f' % (epoch, args.lr)) loss_record = [] # loss1_record = [] # loss2_record = [] # loss3_record = [] local_count = 0 # print('1') for batch_idx, data in enumerate(train_loader): image = data['image'].to(device) depth = data['depth'].to(device) label = data['label'].long().to(device) # print('label', label.shape) output, output_sup1, output_sup2 = model(image, depth) loss1 = loss_func(output, label) loss2 = loss_func(output_sup1, label) loss3 = loss_func(output_sup2, label) loss = loss1 + loss2 + loss3 tq.update(args.batch_size) tq.set_postfix(loss='%.6f' % loss) optimizer.zero_grad() loss.backward() optimizer.step() global_step += 1 local_count += image.data.shape[0] # writer_loss.add_scalar('loss_step', loss, global_step) # writer_loss1.add_scalar('loss1_step', loss1, global_step) # writer_loss2.add_scalar('loss2_step', loss2, global_step) # writer_loss3.add_scalar('loss3_step', loss3, global_step) loss_record.append(loss.item()) # loss1_record.append(loss1.item()) # loss2_record.append(loss2.item()) # loss3_record.append(loss3.item()) if global_step % args.print_freq == 0 or global_step == 1: for name, param in model.named_parameters(): writer_loss.add_histogram(name, param.clone().cpu().data.numpy(), global_step, bins='doane') writer_loss.add_graph(model, [image, depth]) grid_image1 = make_grid(image[:3].clone().cpu().data, 3, normalize=True) writer_loss.add_image('image', grid_image1, global_step) grid_image2 = make_grid(depth[:3].clone().cpu().data, 3, normalize=True) writer_loss.add_image('depth', grid_image2, global_step) grid_image3 = make_grid(utils.color_label( torch.max(output[:3], 1)[1]), 3, normalize=False, range=(0, 255)) writer_loss.add_image('Predicted label', grid_image3, global_step) grid_image4 = make_grid(utils.color_label(label[:3]), 3, normalize=False, range=(0, 255)) writer_loss.add_image('Groundtruth label', grid_image4, global_step) tq.close() loss_train_mean = np.mean(loss_record) with open(log_file, 'a') as f: f.write(str(epoch) + '\t' + str(loss_train_mean)) # loss1_train_mean = np.mean(loss1_record) # loss2_train_mean = np.mean(loss2_record) # loss3_train_mean = np.mean(loss3_record) writer_loss.add_scalar('epoch/loss_epoch_train', float(loss_train_mean), epoch) # writer_loss1.add_scalar('epoch/sub_loss_epoch_train', float(loss1_train_mean), epoch) # writer_loss2.add_scalar('epoch/sub_loss_epoch_train', float(loss2_train_mean), epoch) # writer_loss3.add_scalar('epoch/sub_loss_epoch_train', float(loss3_train_mean), epoch) print('loss for train : %f' % loss_train_mean) print('----validation starting----') # tq_val = tqdm(total=len(val_loader) * args.batch_size) # tq_val.set_description('epoch %d' % epoch) model.eval() val_total_time = 0 with torch.no_grad(): sys.stdout.flush() tbar = tqdm(val_loader) acc_meter = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() a_meter = AverageMeter() b_meter = AverageMeter() for batch_idx, sample in enumerate(tbar): # origin_image = sample['origin_image'].numpy() # origin_depth = sample['origin_depth'].numpy() image_val = sample['image'].to(device) depth_val = sample['depth'].to(device) label_val = sample['label'].numpy() with torch.no_grad(): start = time.time() pred = model(image_val, depth_val) end = time.time() duration = end - start val_total_time += duration # tq_val.set_postfix(fps ='%.4f' % (args.batch_size / (end - start))) print_str = 'Test step [{}/{}].'.format( batch_idx + 1, len(val_loader)) tbar.set_description(print_str) output_val = torch.max(pred, 1)[1] output_val = output_val.squeeze(0).cpu().numpy() acc, pix = accuracy(output_val, label_val) intersection, union = intersectionAndUnion( output_val, label_val, args.num_class) acc_meter.update(acc, pix) a_m, b_m = macc(output_val, label_val, args.num_class) intersection_meter.update(intersection) union_meter.update(union) a_meter.update(a_m) b_meter.update(b_m) fps = len(val_loader) / val_total_time print('fps = %.4f' % fps) tbar.close() mAcc = (a_meter.average() / (b_meter.average() + 1e-10)) with open(log_file, 'a') as f: f.write(' ' + str(mAcc.mean()) + '\n') iou = intersection_meter.sum / (union_meter.sum + 1e-10) writer_acc.add_scalar('epoch/Acc_epoch_train', mAcc.mean(), epoch) print('----validation finished----') model.train() # # 每隔save_epoch_freq个epoch就保存一次权重 if epoch != args.start_epoch: if iou.mean() >= max_miou_val: print('mIoU:', iou.mean()) if not os.path.isdir(args.ckpt_dir): os.mkdir(args.ckpt_dir) save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch, local_count, num_train) max_miou_val = iou.mean() # max_macc_val = mAcc.mean() else: loss_count += 1 torch.cuda.empty_cache()
def inference(): writer_image = SummaryWriter(os.path.join(args.summary_dir, 'segtest')) model = MultiTaskCNN(38, depth_channel=1, pretrained=False, arch='resnet50', use_aspp=False) load_ckpt(model, None, args.last_ckpt, device) model.eval() model = model.to(device) val_data = data_eval.ReadData(transform=torchvision.transforms.Compose( [data_eval.scaleNorm(), data_eval.ToTensor(), Normalize()]), data_dir=args.data_dir) val_loader = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4, pin_memory=False) acc_meter = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() a_meter = AverageMeter() b_meter = AverageMeter() test_total_time = 0 with torch.no_grad(): for batch_idx, sample in enumerate(val_loader): # origin_image = sample['origin_image'].to(device) # origin_depth = sample['origin_depth'].to(device) image = sample['image'].to(device) depth = sample['depth'].to(device) label = sample['label'].numpy() show_label = sample['label'].long().to(device) with torch.no_grad(): time1 = time.time() pred = model(image, depth) time2 = time.time() test_total_time += (time2 - time1) output = torch.max(pred, 1)[1] # # output = output.squeeze(0).cpu().numpy() output = output.cpu().numpy() acc, pix = accuracy(output, label) intersection, union = intersectionAndUnion(output, label, args.num_class) acc_meter.update(acc, pix) a_m, b_m = macc(output, label, args.num_class) intersection_meter.update(intersection) union_meter.update(union) a_meter.update(a_m) b_meter.update(b_m) if batch_idx % 50 == 0: grid_image1 = make_grid(image[:1].clone().cpu().data, 1, normalize=True) writer_image.add_image('image', grid_image1, batch_idx) grid_image2 = make_grid(depth[:1].clone().cpu().data, 1, normalize=True) writer_image.add_image('depth', grid_image2, batch_idx) grid_image3 = make_grid(utils.color_label( torch.max(pred[:1], 1)[1]), 1, normalize=False, range=(0, 255)) writer_image.add_image('Predicted label', grid_image3, batch_idx) grid_image4 = make_grid(utils.color_label(show_label[:1]), 1, normalize=False, range=(0, 255)) writer_image.add_image('Groundtruth label', grid_image4, batch_idx) print('[{}] iter {}, accuracy: {}'.format( datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), batch_idx, acc)) # if batch_idx % 1 == 0: # if args.visualize: # visualize_result(origin_image, origin_depth, label, output, batch_idx, args) # visualize_result(origin_image, origin_depth, label - 1, output - 1, batch_idx, args) print('推理时间:', test_total_time / len(val_data), '\nfps:', len(val_data) / test_total_time) iou = intersection_meter.sum / (union_meter.sum + 1e-10) for i, _iou in enumerate(iou): print('class [{}], IoU: {}'.format(i, _iou)) # mAcc:Prediction和Ground Truth对应位置的“分类”准确率(每个像素) mAcc = (a_meter.average() / (b_meter.average() + 1e-10)) print(mAcc.mean()) print('[Eval Summary]:') print('Mean IoU: {:.4}, Accuracy: {:.2f}%'.format( iou.mean(), acc_meter.average() * 100))
def train(): train_data = ACNet_data.SUNRGBD(transform=transforms.Compose([ACNet_data.scaleNorm(), ACNet_data.RandomScale((1.0, 1.4)), ACNet_data.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), ACNet_data.RandomCrop(image_h, image_w), ACNet_data.RandomFlip(), ACNet_data.ToTensor(), ACNet_data.Normalize()]), phase_train=True, data_dir=args.data_dir) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False) num_train = len(train_data) if args.last_ckpt: model = ACNet_models_V1.ACNet(num_class=40, pretrained=False) else: model = ACNet_models_V1.ACNet(num_class=40, pretrained=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) CEL_weighted = utils.CrossEntropyLoss2d(weight=nyuv2_frq) model.train() model.to(device) CEL_weighted.to(device) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) global_step = 0 if args.last_ckpt: global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device) lr_decay_lambda = lambda epoch: args.lr_decay_rate ** (epoch // args.lr_epoch_per_decay) scheduler = LambdaLR(optimizer, lr_lambda=lr_decay_lambda) writer = SummaryWriter(args.summary_dir) for epoch in range(int(args.start_epoch), args.epochs): scheduler.step(epoch) local_count = 0 last_count = 0 end_time = time.time() if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch: save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch, local_count, num_train) for batch_idx, sample in enumerate(train_loader): image = sample['image'].to(device) depth = sample['depth'].to(device) target_scales = [sample[s].to(device) for s in ['label', 'label2', 'label3', 'label4', 'label5']] optimizer.zero_grad() pred_scales = model(image, depth, args.checkpoint) loss = CEL_weighted(pred_scales, target_scales) loss.backward() optimizer.step() local_count += image.data.shape[0] global_step += 1 if global_step % args.print_freq == 0 or global_step == 1: time_inter = time.time() - end_time count_inter = local_count - last_count print_log(global_step, epoch, local_count, count_inter, num_train, loss, time_inter) end_time = time.time() for name, param in model.named_parameters(): writer.add_histogram(name, param.clone().cpu().data.numpy(), global_step, bins='doane') grid_image = make_grid(image[:3].clone().cpu().data, 3, normalize=True) writer.add_image('image', grid_image, global_step) grid_image = make_grid(depth[:3].clone().cpu().data, 3, normalize=True) writer.add_image('depth', grid_image, global_step) grid_image = make_grid(utils.color_label(torch.max(pred_scales[0][:3], 1)[1] + 1), 3, normalize=False, range=(0, 255)) writer.add_image('Predicted label', grid_image, global_step) grid_image = make_grid(utils.color_label(target_scales[0][:3]), 3, normalize=False, range=(0, 255)) writer.add_image('Groundtruth label', grid_image, global_step) writer.add_scalar('CrossEntropyLoss', loss.data, global_step=global_step) writer.add_scalar('Learning rate', scheduler.get_lr()[0], global_step=global_step) last_count = local_count save_ckpt(args.ckpt_dir, model, optimizer, global_step, args.epochs, 0, num_train) print("Training completed ")
def main(): cfg = Config() # Redirect logs to both console and file. if cfg.log_to_file: ReDirectSTD(cfg.stdout_file, 'stdout', False) ReDirectSTD(cfg.stderr_file, 'stderr', False) # Lazily create SummaryWriter writer = None TVT, TMO = set_devices(cfg.sys_device_ids) if cfg.seed is not None: set_seed(cfg.seed) # Dump the configurations to log. import pprint print('-' * 60) print('cfg.__dict__') pprint.pprint(cfg.__dict__) print('-' * 60) ########### # Dataset # ########### if not cfg.only_test: train_set = create_dataset(**cfg.train_set_kwargs) # The combined dataset does not provide val set currently. val_set = None if cfg.dataset == 'combined' else create_dataset(**cfg.val_set_kwargs) test_sets = [] test_set_names = [] if cfg.dataset == 'combined': for name in ['market1501', 'cuhk03', 'duke']: cfg.test_set_kwargs['name'] = name test_sets.append(create_dataset(**cfg.test_set_kwargs)) test_set_names.append(name) else: test_sets.append(create_dataset(**cfg.test_set_kwargs)) test_set_names.append(cfg.dataset) ########### # Models # ########### if cfg.only_test: model = Model(cfg.net, pretrained=False, last_conv_stride=cfg.last_conv_stride) else: model = Model(cfg.net, path_to_predefined=cfg.net_pretrained_path, last_conv_stride=cfg.last_conv_stride) # This is a ShuffleNet Network. Model(last_conv_stride=cfg.last_conv_stride) ############################# # Criteria and Optimizers # ############################# tri_loss = TripletLoss(margin=cfg.margin) optimizer = optim.Adam(model.parameters(), lr=cfg.base_lr, weight_decay=cfg.weight_decay) #optimizer = optimizers.FusedAdam(model.parameters(), # lr=cfg.base_lr, # weight_decay=cfg.weight_decay) #optimizer = torch.optim.SGD(model.parameters(), cfg.base_lr, # nesterov=True, # momentum=cfg.momentum, # weight_decay=cfg.weight_decay) model.cuda() model, optimizer = amp.initialize(model, optimizer, opt_level=cfg.opt_level, keep_batchnorm_fp32=cfg.keep_batchnorm_fp32, #loss_scale=cfg.loss_scale ) amp.init() # Register function # Bind them together just to save some codes in the following usage. modules_optims = [model, optimizer] # Model wrapper model_w = DataParallel(model) ################################ # May Resume Models and Optims # ################################ if cfg.resume: resume_ep, scores = load_ckpt(modules_optims, cfg.ckpt_file) # May Transfer Models and Optims to Specified Device. Transferring optimizer # is to cope with the case when you load the checkpoint to a new device. TMO(modules_optims) ######## # Test # ######## def test(load_model_weight=False): if load_model_weight: if cfg.model_weight_file != '': map_location = (lambda storage, loc: storage) sd = torch.load(cfg.model_weight_file, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format(cfg.model_weight_file)) else: load_ckpt(modules_optims, cfg.ckpt_file) for test_set, name in zip(test_sets, test_set_names): feature_map = ExtractFeature(model_w, TVT) test_set.set_feat_func(feature_map) print('\n=========> Test on dataset: {} <=========\n'.format(name)) test_set.eval( normalize_feat=cfg.normalize_feature, verbose=True) def validate(): if val_set.extract_feat_func is None: feature_map = ExtractFeature(model_w, TVT) val_set.set_feat_func(feature_map) print('\n=========> Test on validation set <=========\n') mAP, cmc_scores, _, _ = val_set.eval( normalize_feat=cfg.normalize_feature, to_re_rank=False, verbose=False) print() return mAP, cmc_scores[0] if cfg.only_test: test(load_model_weight=True) return ############ # Training # ############ start_ep = resume_ep if cfg.resume else 0 for ep in range(start_ep, cfg.total_epochs): # Adjust Learning Rate if cfg.lr_decay_type == 'exp': adjust_lr_exp( optimizer, cfg.base_lr, ep + 1, cfg.total_epochs, cfg.exp_decay_at_epoch) else: adjust_lr_staircase( optimizer, cfg.base_lr, ep + 1, cfg.staircase_decay_at_epochs, cfg.staircase_decay_multiply_factor) may_set_mode(modules_optims, 'train') # For recording precision, satisfying margin, etc prec_meter = AverageMeter() sm_meter = AverageMeter() dist_ap_meter = AverageMeter() dist_an_meter = AverageMeter() loss_meter = AverageMeter() ep_st = time.time() step = 0 epoch_done = False while not epoch_done: step += 1 step_st = time.time() ims, im_names, labels, mirrored, epoch_done = train_set.next_batch() ims_var = Variable(TVT(torch.from_numpy(ims).float())) labels_t = TVT(torch.from_numpy(labels).long()) feat = model_w(ims_var) loss, p_inds, n_inds, dist_ap, dist_an, dist_mat = global_loss( tri_loss, feat, labels_t, normalize_feature=cfg.normalize_feature) optimizer.zero_grad() with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() optimizer.step() ############ # Step Log # ############ # precision prec = (dist_an > dist_ap).data.float().mean() # the proportion of triplets that satisfy margin sm = (dist_an > dist_ap + cfg.margin).data.float().mean() # average (anchor, positive) distance d_ap = dist_ap.data.mean() # average (anchor, negative) distance d_an = dist_an.data.mean() prec_meter.update(prec) sm_meter.update(sm) dist_ap_meter.update(d_ap) dist_an_meter.update(d_an) loss_meter.update(to_scalar(loss)) if step % cfg.steps_per_log == 0: time_log = '\tStep {}/Ep {}, {:.2f}s'.format( step, ep + 1, time.time() - step_st, ) tri_log = (', prec {:.2%}, sm {:.2%}, ' 'd_ap {:.4f}, d_an {:.4f}, ' 'loss {:.4f}'.format( prec_meter.val, sm_meter.val, dist_ap_meter.val, dist_an_meter.val, loss_meter.val, )) log = time_log + tri_log print(log) ############# # Epoch Log # ############# time_log = 'Ep {}, {:.2f}s'.format(ep + 1, time.time() - ep_st) tri_log = (', prec {:.2%}, sm {:.2%}, ' 'd_ap {:.4f}, d_an {:.4f}, ' 'loss {:.4f}'.format( prec_meter.avg, sm_meter.avg, dist_ap_meter.avg, dist_an_meter.avg, loss_meter.avg, )) log = time_log + tri_log print(log) ########################## # Test on Validation Set # ########################## mAP, Rank1 = 0, 0 if ((ep + 1) % cfg.epochs_per_val == 0) and (val_set is not None): mAP, Rank1 = validate() # Log to TensorBoard if cfg.log_to_file: if writer is None: writer = SummaryWriter(log_dir=osp.join(cfg.exp_dir, 'tensorboard')) writer.add_scalars( 'val scores', dict(mAP=mAP, Rank1=Rank1), ep) writer.add_scalars( 'loss', dict(loss=loss_meter.avg, ), ep) writer.add_scalars( 'precision', dict(precision=prec_meter.avg, ), ep) writer.add_scalars( 'satisfy_margin', dict(satisfy_margin=sm_meter.avg, ), ep) writer.add_scalars( 'average_distance', dict(dist_ap=dist_ap_meter.avg, dist_an=dist_an_meter.avg, ), ep) # save ckpt if cfg.log_to_file: save_ckpt(modules_optims, ep + 1, 0, cfg.ckpt_file) ######## # Test # ######## test(load_model_weight=False)