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)
