def step(args, split, epoch, loader, model, optimizer = None, M = None, f = None, tag = None):
losses, mpjpe, mpjpe_r = AverageMeter(), AverageMeter(), AverageMeter()
viewLosses, shapeLosses, supLosses = AverageMeter(), AverageMeter(), AverageMeter()
if split == 'train':
model.train()
else:
model.eval()
bar = Bar('{}'.format(ref.category), max=len(loader))
nViews = loader.dataset.nViews
for i, (input, target, meta) in enumerate(loader):
input_var = torch.autograd.Variable(input)
target_var = torch.autograd.Variable(target)
output = model(input_var)
loss = ShapeConsistencyCriterion(nViews, supWeight = 1, unSupWeight = args.shapeWeight, M = M)(output, target_var, torch.autograd.Variable(meta))
if split == 'test':
for j in range(input.numpy().shape[0]):
img = (input.numpy()[j] * 255).transpose(1, 2, 0).astype(np.uint8)
cv2.imwrite('{}/img_{}/{}.png'.format(args.save_path, tag, i * input.numpy().shape[0] + j), img)
gt = target.cpu().numpy()[j]
pred = (output.data).cpu().numpy()[j]
vis = meta.cpu().numpy()[j][5:]
for t in range(ref.J):
f.write('{} {} {} '.format(pred[t * 3], pred[t * 3 + 1], pred[t * 3 + 2]))
f.write('\n')
for t in range(ref.J):
f.write('{} {} {} '.format(gt[t, 0], gt[t, 1], gt[t, 2]))
f.write('\n')
if args.saveVis:
for t in range(ref.J):
f.write('{} 0 0 '.format(vis[t]))
f.write('\n')
mpjpe_this = accuracy(output.data, target, meta)
mpjpe_r_this = accuracy_dis(output.data, target, meta)
shapeLoss = shapeConsistency(output.data, meta, nViews, M, split = split)
losses.update(loss.data[0], input.size(0))
shapeLosses.update(shapeLoss, input.size(0))
mpjpe.update(mpjpe_this, input.size(0))
mpjpe_r.update(mpjpe_r_this, input.size(0))
if split == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
Bar.suffix = '{split:10}: [{0:2}][{1:3}/{2:3}] | Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | shapeLoss {shapeLoss.avg:.6f} | AE {mpjpe.avg:.6f} | ShapeDis {mpjpe_r.avg:.6f}'.format(epoch, i, len(loader), total=bar.elapsed_td, eta=bar.eta_td, loss=losses, mpjpe=mpjpe, split = split, shapeLoss = shapeLosses, mpjpe_r = mpjpe_r)
bar.next()
bar.finish()
return mpjpe.avg, losses.avg, shapeLosses.avg
def run_epoch(self, phase, epoch, data_loader):
model_with_loss = self.model_with_loss
if len(self.opt.gpus) > 1:
if self.opt.weight_strategy == 'GRADNORM':
model_with_loss.module.loss.optimizer = get_loss_optimizer(model=model_with_loss.module.loss.loss_model,
opt=self.opt)
model_with_loss.module.loss.update_weight(epoch)
else:
if self.opt.weight_strategy == 'GRADNORM':
model_with_loss.loss.optimizer = get_loss_optimizer(model=model_with_loss.loss.loss_model, opt=self.opt)
model_with_loss.loss.update_weight(epoch)
if phase == 'train':
model_with_loss.train()
else:
if len(self.opt.gpus) > 1:
model_with_loss = self.model_with_loss.module
model_with_loss.eval()
torch.cuda.empty_cache()
opt = self.opt
results = {}
data_time, batch_time = AverageMeter(), AverageMeter()
avg_loss_stats = {l: AverageMeter() for l in self.loss_stats} #\
# if l == 'tot' or opt.weights[l] > 0}
num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
bar = Bar('{}/{}'.format(opt.task, opt.exp_id), max=num_iters)
end = time.time()
for iter_id, batch in enumerate(data_loader):
if iter_id >= num_iters:
break
data_time.update(time.time() - end)
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=opt.device, non_blocking=True)
output, loss, loss_stats = model_with_loss(batch, epoch)
loss = loss.mean()
if phase == 'train':
self.optimizer.zero_grad()
if opt.weight_strategy == 'GRADNORM':
loss.backward(retain_graph=True)
if len(self.opt.gpus) > 1:
# model_with_loss.module.loss.loss_model.update_weight(model_with_loss.module.model, model_with_loss.module.loss.optimizer, loss_stats)
# torch.sum(loss_stats['update']).backward()
model_with_loss.module.loss.optimizer.zero_grad()
temp_grad = torch.autograd.grad(torch.sum(loss_stats['update']),
model_with_loss.module.loss.loss_model.weight)[0]
grad_norm = torch.norm(temp_grad.data, 1)
print(grad_norm)
if grad_norm > opt.gradnorm_thred:
temp_grad = torch.zeros_like(temp_grad)
model_with_loss.module.loss.loss_model.weight.grad = temp_grad
model_with_loss.module.loss.optimizer.step()
else:
# model_with_loss.loss.loss_model.update_weight(model_with_loss.model,
# model_with_loss.loss.optimizer, loss_stats)
model_with_loss.loss.optimizer.zero_grad()
temp_grad = torch.autograd.grad(loss_stats['update'],
model_with_loss.loss.loss_model.
weight)[0]
grad_norm = torch.norm(temp_grad.data, 1)
if grad_norm > opt.gradnorm_thred:
temp_grad = torch.zeros_like(temp_grad)
model_with_loss.loss.loss_model.weight.grad = temp_grad
model_with_loss.loss.optimizer.step()
else:
# torch.autograd.grad(loss, model_with_loss.model.padnet.parameters())
loss.backward()
self.optimizer.step()
if opt.weight_strategy == 'UNCER':
if len(self.opt.gpus) > 1:
model_with_loss.module.loss.optimizer.step()
model_with_loss.module.loss.optimizer.zero_grad()
print(model_with_loss.module.loss.group_weight)
else:
model_with_loss.loss.optimizer.step()
model_with_loss.loss.optimizer.zero_grad()
print(model_with_loss.loss.group_weight)
batch_time.update(time.time() - end)
end = time.time()
Bar.suffix = '{phase}: [{0}][{1}/{2}]|Tot: {total:} |ETA: {eta:} '.format(
epoch, iter_id, num_iters, phase=phase,
total=bar.elapsed_td, eta=bar.eta_td)
for l in avg_loss_stats:
avg_loss_stats[l].update(
loss_stats[l].mean().item(), batch['image'].size(0))
Bar.suffix = Bar.suffix + '|{} {:.4f} '.format(l, avg_loss_stats[l].avg)
Bar.suffix = Bar.suffix + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
'|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
if opt.print_iter > 0: # If not using progress bar
if iter_id % opt.print_iter == 0:
print('{}/{}| {}'.format(opt.task, opt.exp_id, Bar.suffix))
else:
bar.next()
if opt.debug > 0:
self.debug(batch, output, iter_id, dataset=data_loader.dataset)
del output, loss, loss_stats
bar.finish()
ret = {k: v.avg for k, v in avg_loss_stats.items()}
ret['time'] = bar.elapsed_td.total_seconds() / 60.
if len(self.opt.gpus) > 1:
model_with_loss.module.loss.update_loss(epoch, ret)
else:
model_with_loss.loss.update_loss(epoch, ret)
return ret, results
def run_epoch(self, phase, epoch, data_loader):
model_with_loss = self.model_with_loss
if phase == 'train':
model_with_loss.train()
else:
if len(self.opt.gpus) > 1:
model_with_loss = self.model_with_loss.module
model_with_loss.eval()
torch.cuda.empty_cache()
opt = self.opt
results = {}
data_time, batch_time = AverageMeter(), AverageMeter()
avg_loss_stats = {l: AverageMeter() for l in self.loss_stats}
num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
bar = Bar('{}/{}'.format(opt.task, opt.exp_id), max=num_iters)
end = time.time()
for iter_id, batch in enumerate(data_loader):
if iter_id >= num_iters:
break
data_time.update(time.time() - end)
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=opt.device,
non_blocking=True)
output, loss, loss_stats = model_with_loss(batch)
loss = loss.mean()
if phase == 'train':
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
Bar.suffix = '{phase}: [{0}][{1}/{2}]|Tot: {total:} |ETA: {eta:} '.format(
epoch,
iter_id,
num_iters,
phase=phase,
total=bar.elapsed_td,
eta=bar.eta_td)
for l in avg_loss_stats:
avg_loss_stats[l].update(loss_stats[l].mean().item(),
batch['input'].size(0))
Bar.suffix = Bar.suffix + '|{} {:.4f} '.format(
l, avg_loss_stats[l].avg)
if not opt.hide_data_time:
Bar.suffix = Bar.suffix + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
'|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
if opt.print_iter > 0:
if iter_id % opt.print_iter == 0:
print('{}/{}| {}'.format(opt.task, opt.exp_id, Bar.suffix))
else:
bar.next()
if opt.debug > 0:
self.debug(batch, output, iter_id)
if opt.test:
self.save_result(output, batch, results)
del output, loss, loss_stats
bar.finish()
ret = {k: v.avg for k, v in avg_loss_stats.items()}
ret['time'] = bar.elapsed_td.total_seconds() / 60.
return ret, results
def validate(val_loader, net, criterion, curr_epoch):
# the following code is written assuming that batch size is 1
net.eval()
if args['gpu']:
torch.cuda.empty_cache()
start = time.time()
val_loss = AverageMeter()
acc_meter = AverageMeter()
fwIoU_meter = AverageMeter()
for vi, (imgs, labels) in enumerate(val_loader):
imgs = imgs.float()
labels = labels.long()
if args['gpu']:
imgs = imgs.cuda().float()
labels = labels.cuda().long()
with torch.no_grad():
outputs, aux = net(imgs)
loss = criterion(outputs, labels)
val_loss.update(loss.cpu().detach().numpy())
outputs = outputs.cpu().detach()
labels = labels.cpu().detach().numpy()
_, preds = torch.max(outputs, dim=1)
preds = preds.numpy()
for (pred, label) in zip(preds, labels):
acc, valid_sum = accuracy(pred, label)
fwiou = FWIoU(pred.squeeze(), label.squeeze(), ignore_zero=True)
acc_meter.update(acc)
fwIoU_meter.update(fwiou)
if curr_epoch % args['predict_step'] == 0 and vi == 0:
pred_color = RS.Index2Color(preds[0])
io.imsave(os.path.join(args['pred_dir'], NET_NAME + '.png'),
pred_color)
print('Prediction saved!')
curr_time = time.time() - start
print('%.1fs Val loss: %.2f, Accuracy: %.2f, fwIoU: %.2f' %
(curr_time, val_loss.average(), acc_meter.average() * 100,
fwIoU_meter.average() * 100))
writer.add_scalar('val_loss', val_loss.average(), curr_epoch)
writer.add_scalar('val_Accuracy', acc_meter.average(), curr_epoch)
writer.add_scalar('val_fwIoU', fwIoU_meter.average(), curr_epoch)
return acc_meter.avg, fwIoU_meter.avg, val_loss.avg
def validate(args, val_dataloader, model, auxiliarynet, epoch):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
error = AverageMeter('error', ':6.2f')
progress = ProgressMeter(len(val_dataloader),
batch_time,
data_time,
losses,
error,
prefix="Val Epoch: [{}]".format(epoch))
model.eval()
# auxiliarynet.eval()
end = time.time()
with torch.no_grad():
end = time.time()
for i, (patch, gaze_norm_g, head_norm,
rot_vec_norm) in enumerate(val_dataloader):
# measure data loading time
data_time.update(time.time() - end)
patch = patch.to(args.device)
gaze_norm_g = gaze_norm_g.to(args.device)
head_norm = head_norm.to(args.device)
rot_vec_norm = rot_vec_norm.to(args.device)
# model = model.to(args.device)
gaze_pred, _ = model(patch)
# hp_pred = auxiliarynet(features)
head_norm = 10 * head_norm
gaze_norm_g = 10 * gaze_norm_g
# loss = criterion(gaze_norm_g, head_norm, gaze_pred[:,0:2], gaze_pred[:,2:4])
angle_error = mean_angle_error(
gaze_pred.cpu().detach().numpy() / 10,
gaze_norm_g.cpu().detach().numpy() / 10,
rot_vec_norm.cpu().detach().numpy())
# losses.update(loss.item())
error.update(angle_error)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if (i + 1) % args.print_freq == 0:
progress.print(i + 1)
# img = patch.cpu().detach().numpy()[0].deepcopy()
# to_visualize = draw_gaze(img[0], (0.25 * img.shape[1], 0.25 * img.shape[1]), gaze_pred,
# gaze_norm_g, length=80.0, thickness=1)
# cv2.imshow('vis', to_visualize)
# cv2.waitKey(1)
return losses.get_avg(), error.get_avg()
def validate(args, val_loader, model, epoch, win_feats5, win_fusion, viz,
global_step):
batch_time = AverageMeter()
data_time = AverageMeter()
feats5_losses = AverageMeter()
fusion_losses = AverageMeter()
total_losses = AverageMeter()
# switch to train mode
model.eval()
torch.no_grad()
end = time.time()
for i, (img, target) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
# Input for Image CNN.
img_var = utils.check_gpu(0, img) # BS X 3 X H X W
target_var = utils.check_gpu(0, target) # BS X H X W X NUM_CLASSES
bs = img.size()[0]
score_feats5, fused_feats = model(
img_var) # BS X NUM_CLASSES X 472 X 472
feats5_loss = WeightedMultiLabelSigmoidLoss(score_feats5, target_var)
fused_feats_loss = WeightedMultiLabelSigmoidLoss(
fused_feats, target_var)
loss = feats5_loss + fused_feats_loss
feats5_losses.update(feats5_loss.data, bs)
fusion_losses.update(fused_feats_loss.data, bs)
total_losses.update(loss.data, bs)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# clear memory
del img_var
del target_var
del score_feats5
del fused_feats_loss
del feats5_loss
torch.cuda.empty_cache()
if (i % args.print_freq == 0):
print("\n\n")
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Total Loss {total_loss.val:.4f} ({total_loss.avg:.4f})\n'.
format(epoch,
i,
len(val_loader),
batch_time=batch_time,
data_time=data_time,
total_loss=total_losses))
#viz.line(win=win_feats5, name='val_feats5', update='append', X=np.array([global_step]), Y=np.array([feats5_losses.avg]))
#viz.line(win=win_fusion, name='val_fusion', update='append', X=np.array([global_step]), Y=np.array([fusion_losses.avg]))
return fusion_losses.avg
def train(train_loader, miml, encoder, decoder, criterion, encoder_optimizer,
decoder_optimizer, epoch, writer):
"""
Performs one epoch's training.
:param train_loader: DataLoader for training data
:param encoder: encoder model
:param decoder: decoder model
:param criterion: loss layer
:param encoder_optimizer: optimizer to update encoder's weights (if fine-tuning)
:param decoder_optimizer: optimizer to update decoder's weights
:param epoch: epoch number
"""
encoder.train()
decoder.train() # train mode (dropout and batchnorm is used)
miml.train()
total_step = len(train_loader)
batch_time = AverageMeter() # forward prop. + back prop. time
data_time = AverageMeter() # data loading time
losses = AverageMeter() # loss (per word decoded)
top5accs = AverageMeter() # top5 accuracy
start = time.time()
# Batches
for i, (imgs, caps, caplens) in enumerate(train_loader):
data_time.update(time.time() - start)
# Move to GPU, if available
imgs = imgs.to(device)
caps = caps.to(device)
caplens = caplens.to(device)
# Forward prop.
# attrs = miml(imgs)
# imgs = encoder(imgs)
scores, caps_sorted, decode_lengths, alphas, sort_ind = decoder(
miml(imgs), encoder(imgs), caps, caplens)
# Since we decoded starting with <start>, the targets are all words after <start>, up to <end>
targets = caps_sorted[:, 1:]
# Remove timesteps that we didn't decode at, or are pads
# pack_padded_sequence is an easy trick to do this
# torch在计算时会自动除去pad，这样不带pad计算不影响精度
scores, _ = pack_padded_sequence(scores,
decode_lengths,
batch_first=True)
targets, _ = pack_padded_sequence(targets,
decode_lengths,
batch_first=True)
# Calculate loss
loss = criterion(scores, targets)
loss += alpha_c * ((1. - alphas.sum(dim=1))**2).mean()
# Back prop.
decoder_optimizer.zero_grad()
if encoder_optimizer is not None:
encoder_optimizer.zero_grad()
loss.backward()
# Clip gradients
if grad_clip is not None:
clip_gradient(decoder_optimizer, grad_clip)
if encoder_optimizer is not None:
clip_gradient(encoder_optimizer, grad_clip)
# Update weights
if encoder_optimizer is not None:
encoder_optimizer.step()
decoder_optimizer.step()
# Keep track of metrics
top5 = accuracy(scores, targets, 5)
losses.update(loss.item(), sum(decode_lengths))
top5accs.update(top5, sum(decode_lengths))
batch_time.update(time.time() - start)
start = time.time()
# Print status
if i % print_freq == 0:
writer.add_scalars('train', {
'loss': loss.item(),
'mAp': top5accs.val
}, epoch * total_step + i)
print('Epoch: [{0}][{1}/{2}]\t'
'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data Load Time {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Top-5 Accuracy {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch,
i,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=losses,
top5=top5accs))
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=20,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
# Create model
model = Model().cuda()
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
prefix, args.lr, args.epochs, args.bs),
flush_secs=30)
# Loss
l1_criterion = nn.L1Loss()
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 300 == 0:
LogProgress(model, writer, test_loader, niter)
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
def update_network(self, **kwargs):
stop = False
self.train_data['source']['iterator'] = iter(
self.train_data['source']['loader'])
self.train_data['target']['iterator'] = iter(
self.train_data['target']['loader'])
self.iters_per_epoch = len(self.train_data['target']['loader'])
iters_counter_within_epoch = 0
data_time = AverageMeter()
batch_time = AverageMeter()
classifier_loss = AverageMeter()
feature_extractor_loss = AverageMeter()
prec1_fs = AverageMeter()
prec1_ft = AverageMeter()
self.feature_extractor.train()
self.classifier.train()
end = time.time()
if self.opt.TRAIN.PROCESS_COUNTER == 'epoch':
lam = 2 / (1 + math.exp(
-1 * 10 * self.epoch / self.opt.TRAIN.MAX_EPOCH)) - 1
self.update_lr()
print('value of lam is: %3f' % (lam))
while not stop:
if self.opt.TRAIN.PROCESS_COUNTER == 'iteration':
lam = 2 / (1 + math.exp(
-1 * 10 * self.iters /
(self.opt.TRAIN.MAX_EPOCH * self.iters_per_epoch))) - 1
print('value of lam is: %3f' % (lam))
self.update_lr()
source_data, source_gt = self.get_samples('source')
target_data, _ = self.get_samples('target')
source_data = to_cuda(source_data)
source_gt = to_cuda(source_gt)
target_data = to_cuda(target_data)
data_time.update(time.time() - end)
feature_source = self.feature_extractor(source_data)
output_source = self.classifier(feature_source)
feature_target = self.feature_extractor(target_data)
output_target = self.classifier(feature_target)
loss_task_fs = self.CELoss(output_source[:, :self.num_classes],
source_gt)
loss_task_ft = self.CELoss(output_source[:, self.num_classes:],
source_gt)
loss_discrim_source = self.SourceDiscrimLoss(output_source)
loss_discrim_target = self.TargetDiscrimLoss(output_target)
loss_summary_classifier = loss_task_fs + loss_task_ft + loss_discrim_source + loss_discrim_target
source_gt_for_ft_in_fst = source_gt + self.num_classes
loss_confusion_source = 0.5 * self.CELoss(
output_source, source_gt) + 0.5 * self.CELoss(
output_source, source_gt_for_ft_in_fst)
loss_confusion_target = 0.5 * self.SourceDiscrimLoss(
output_target) + 0.5 * self.TargetDiscrimLoss(output_target)
loss_em = self.ConcatenatedEMLoss(output_target)
loss_summary_feature_extractor = loss_confusion_source + lam * (
loss_confusion_target + loss_em)
self.optimizer_classifier.zero_grad()
loss_summary_classifier.backward(retain_graph=True)
self.optimizer_classifier.step()
self.optimizer_feature_extractor.zero_grad()
loss_summary_feature_extractor.backward()
self.optimizer_feature_extractor.step()
classifier_loss.update(loss_summary_classifier,
source_data.size()[0])
feature_extractor_loss.update(loss_summary_feature_extractor,
source_data.size()[0])
prec1_fs.update(
accuracy(output_source[:, :self.num_classes], source_gt),
source_data.size()[0])
prec1_ft.update(
accuracy(output_source[:, self.num_classes:], source_gt),
source_data.size()[0])
print(" Train:epoch: %d:[%d/%d], LossCla: %3f, LossFeat: %3f, AccFs: %3f, AccFt: %3f" % \
(self.epoch, iters_counter_within_epoch, self.iters_per_epoch, classifier_loss.avg, feature_extractor_loss.avg, prec1_fs.avg, prec1_ft.avg))
batch_time.update(time.time() - end)
end = time.time()
self.iters += 1
iters_counter_within_epoch += 1
if iters_counter_within_epoch >= self.iters_per_epoch:
log = open(os.path.join(self.opt.SAVE_DIR, 'log.txt'), 'a')
log.write("\n")
log.write(" Train:epoch: %d:[%d/%d], LossCla: %3f, LossFeat: %3f, AccFs: %3f, AccFt: %3f" % \
(self.epoch, iters_counter_within_epoch, self.iters_per_epoch, classifier_loss.avg, feature_extractor_loss.avg, prec1_fs.avg, prec1_ft.avg))
log.close()
stop = True
def test(self):
self.feature_extractor.eval()
self.classifier.eval()
prec1_fs = AverageMeter()
prec1_ft = AverageMeter()
counter_all_fs = torch.FloatTensor(
self.opt.DATASET.NUM_CLASSES).fill_(0)
counter_all_ft = torch.FloatTensor(
self.opt.DATASET.NUM_CLASSES).fill_(0)
counter_acc_fs = torch.FloatTensor(
self.opt.DATASET.NUM_CLASSES).fill_(0)
counter_acc_ft = torch.FloatTensor(
self.opt.DATASET.NUM_CLASSES).fill_(0)
for i, (input, target) in enumerate(self.test_data['loader']):
input, target = to_cuda(input), to_cuda(target)
with torch.no_grad():
feature_test = self.feature_extractor(input)
output_test = self.classifier(feature_test)
if self.opt.EVAL_METRIC == 'accu':
prec1_fs_iter = accuracy(output_test[:, :self.num_classes],
target)
prec1_ft_iter = accuracy(output_test[:, self.num_classes:],
target)
prec1_fs.update(prec1_fs_iter, input.size(0))
prec1_ft.update(prec1_ft_iter, input.size(0))
if i % self.opt.PRINT_STEP == 0:
print(" Test:epoch: %d:[%d/%d], AccFs: %3f, AccFt: %3f" % \
(self.epoch, i, len(self.test_data['loader']), prec1_fs.avg, prec1_ft.avg))
elif self.opt.EVAL_METRIC == 'accu_mean':
prec1_ft_iter = accuracy(output_test[:, self.num_classes:],
target)
prec1_ft.update(prec1_ft_iter, input.size(0))
counter_all_fs, counter_acc_fs = accuracy_for_each_class(
output_test[:, :self.num_classes], target, counter_all_fs,
counter_acc_fs)
counter_all_ft, counter_acc_ft = accuracy_for_each_class(
output_test[:, self.num_classes:], target, counter_all_ft,
counter_acc_ft)
if i % self.opt.PRINT_STEP == 0:
print(" Test:epoch: %d:[%d/%d], Task: %3f" % \
(self.epoch, i, len(self.test_data['loader']), prec1_ft.avg))
else:
raise NotImplementedError
acc_for_each_class_fs = counter_acc_fs / counter_all_fs
acc_for_each_class_ft = counter_acc_ft / counter_all_ft
log = open(os.path.join(self.opt.SAVE_DIR, 'log.txt'), 'a')
log.write("\n")
if self.opt.EVAL_METRIC == 'accu':
log.write(
" Test:epoch: %d, AccFs: %3f, AccFt: %3f" % \
(self.epoch, prec1_fs.avg, prec1_ft.avg))
log.close()
return max(prec1_fs.avg, prec1_ft.avg)
elif self.opt.EVAL_METRIC == 'accu_mean':
log.write(
" Test:epoch: %d, AccFs: %3f, AccFt: %3f" % \
(self.epoch,acc_for_each_class_fs.mean(), acc_for_each_class_ft.mean()))
log.write(
"\nClass-wise Acc of Ft:") ## based on the task classifier.
for i in range(self.opt.DATASET.NUM_CLASSES):
if i == 0:
log.write("%dst: %3f" % (i + 1, acc_for_each_class_ft[i]))
elif i == 1:
log.write(", %dnd: %3f" %
(i + 1, acc_for_each_class_ft[i]))
elif i == 2:
log.write(", %drd: %3f" %
(i + 1, acc_for_each_class_ft[i]))
else:
log.write(", %dth: %3f" %
(i + 1, acc_for_each_class_ft[i]))
log.close()
return max(acc_for_each_class_ft.mean(),
acc_for_each_class_fs.mean())
def test_gallery(net, dataloader, output_dir, thresh=0.):
"""test gallery images"""
with open('config.yml', 'r') as f:
config = yaml.load(f)
num_images = len(dataloader.dataset)
all_boxes = []
all_features = []
end = time.time()
time_cost = AverageMeter()
net.eval()
for i, data in enumerate(dataloader):
with torch.no_grad():
im, (orig_shape, im_info) = data
im = im.to(device)
im_info = im_info.numpy().squeeze(0)
orig_shape = [x.item() for x in orig_shape]
scores, bbox_pred, rois, features = net.forward(im, None, im_info)
boxes = rois[:, 1:5] / im_info[2]
scores = np.reshape(scores, [scores.shape[0], -1])
bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
if config['test_bbox_reg']:
# Apply bounding-box regression deltas
box_deltas = bbox_pred
pred_boxes = bbox_transform_inv(
torch.from_numpy(boxes), torch.from_numpy(box_deltas)).numpy()
pred_boxes = clip_boxes(pred_boxes, orig_shape)
else:
# Simply repeat the boxes, once for each class
pred_boxes = np.tile(boxes, (1, scores.shape[1]))
boxes = pred_boxes
# skip j = 0, because it's the background class
j = 1
inds = np.where(scores[:, j] > thresh)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
cls_dets = np.hstack(
(cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32,
copy=False)
keep = nms(torch.from_numpy(cls_dets),
config['test_nms']).numpy() if cls_dets.size > 0 else []
cls_dets = cls_dets[keep, :]
all_boxes.append(cls_dets)
all_features.append(features[inds][keep])
time_cost.update(time.time() - end)
end = time.time()
print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, num_images,
time_cost.avg))
det_file = os.path.join(output_dir, 'gboxes.pkl')
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
feature_file = os.path.join(output_dir, 'gfeatures.pkl')
with open(feature_file, 'wb') as f:
pickle.dump(all_features, f, pickle.HIGHEST_PROTOCOL)
return all_boxes, all_features
def scan_train(train_loader,
model,
criterion,
optimizer,
epoch,
update_cluster_head_only=False):
"""
Train w/ SCAN-Loss
"""
total_losses = AverageMeter('Total Loss', ':.4e')
consistency_losses = AverageMeter('Consistency Loss', ':.4e')
entropy_losses = AverageMeter('Entropy', ':.4e')
progress = ProgressMeter(
len(train_loader), [total_losses, consistency_losses, entropy_losses],
prefix="Epoch: [{}]".format(epoch))
if update_cluster_head_only:
model.eval() # No need to update BN
else:
model.train() # Update BN
for i, batch in enumerate(train_loader):
# Forward pass
anchors = batch['anchor'].cuda(non_blocking=True)
neighbors = batch['neighbor'].cuda(non_blocking=True)
if update_cluster_head_only: # Only calculate gradient for backprop of linear layer
with torch.no_grad():
anchors_features = model(anchors, forward_pass='******')
neighbors_features = model(neighbors, forward_pass='******')
anchors_output = model(anchors_features, forward_pass='******')
neighbors_output = model(neighbors_features, forward_pass='******')
else: # Calculate gradient for backprop of complete network
anchors_output = model(anchors)
neighbors_output = model(neighbors)
# Loss for every head
total_loss, consistency_loss, entropy_loss = [], [], []
for anchors_output_subhead, neighbors_output_subhead in zip(
anchors_output, neighbors_output):
total_loss_, consistency_loss_, entropy_loss_ = criterion(
anchors_output_subhead, neighbors_output_subhead)
total_loss.append(total_loss_)
consistency_loss.append(consistency_loss_)
entropy_loss.append(entropy_loss_)
# Register the mean loss and backprop the total loss to cover all subheads
total_losses.update(np.mean([v.item() for v in total_loss]))
consistency_losses.update(np.mean([v.item()
for v in consistency_loss]))
entropy_losses.update(np.mean([v.item() for v in entropy_loss]))
total_loss = torch.sum(torch.stack(total_loss, dim=0))
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if i % 25 == 0:
progress.display(i)
def test(opt, model, loader):
"""
Validate the model at the current state
Args:
opt (Namspace): training options
model (LaneNet): a LaneNet model
criterion: a CrossEntropyLoss criterion
loader: val data loader
Returns:
The average loss value on val data
"""
model.eval()
run_time = AverageMeter()
end = time.time()
palette = get_palette(opt.dataset, opt.num_classes + 1)
total_inter, total_union, total_correct, total_label = 0, 0, 0, 0
pixAcc = 0
IoU = 0
mIoU = 0
pbar = tqdm(loader)
with torch.no_grad():
for data in pbar:
images, labels, sizes, image_names, org_images = data
sizes = sizes[0].numpy()
images = Variable(images)
N_, C_, H_, W_ = images.shape
if torch.cuda.is_available():
images = images.cuda()
preds = model(images)
preds = gather(preds, 0, dim=0)
if 'dsn' in opt.model_type:
preds = preds[-1]
preds = F.upsample(input=preds,
size=(H_, W_),
mode='bilinear',
align_corners=True)
output = preds.cpu().data.numpy().transpose(0, 2, 3, 1)
seg_pred = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)
# store images
if opt.store_output:
for i in range(N_):
output_im = Image.fromarray(seg_pred[i])
output_im.putpalette(palette)
output_file = os.path.join(opt.output_dir,
image_names[i] + '.png')
output_im.save(output_file)
src_img = org_images[i].data.numpy()
if opt.dataset == 'cityscapes':
drivable_img = np.where(seg_pred[i] == 0, 0,
19).astype(np.uint8)
drivable_img = Image.fromarray(drivable_img)
drivable_img.putpalette(palette)
drivable_img = np.array(
drivable_img.convert('RGB')).astype(src_img.dtype)
#overlay_img = cv2.addWeighted(src_img, 1.0, drivable_img, 1.0, 0)
src_img[drivable_img > 0] = 0
else:
drivable_img = seg_pred[i]
drivable_img = Image.fromarray(drivable_img)
drivable_img.putpalette(palette)
drivable_img = np.array(
drivable_img.convert('RGB')).astype(src_img.dtype)
overlay_img = cv2.add(src_img, drivable_img)
output_file = os.path.join(
opt.output_dir, image_names[i] + '_drivable.png')
cv2.imwrite(output_file,
cv2.cvtColor(overlay_img, cv2.COLOR_RGB2BGR))
if len(labels) > 0:
labels[labels == opt.ignore_label] = -1 #
correct, labeled = batch_pix_accuracy(preds.data, labels)
inter, union = batch_intersection_union(
preds.data, labels, opt.num_classes)
total_correct += correct
total_label += labeled
total_inter += inter
total_union += union
pixAcc = 1.0 * total_correct / (np.spacing(1) + total_label)
IoU = 1.0 * total_inter / (np.spacing(1) + total_union)
mIoU = IoU.mean()
# measure speed test
run_time.update(time.time() - end)
nd = time.time()
fps = N_ / run_time.avg
pbar.set_description(
'Average run time: {:.3f} fps, pixAcc={:.6f}, mIoU={:.6f}'.
format(fps, pixAcc, mIoU))
print({'meanIU': mIoU, 'IU_array': IoU})
return mIoU
def train_fixed(starting_epoch, data_loader, fixed_cnn, criterion, fixed_cnn_optmizer, fixed_cnn_scheduler, utilHelper):
global GLOBAL_STEP, reduction_arc, cell_arc, EVAL_BEST_ACC, EVAL_STEP, TRAIN_HISTORY
for epoch in tqdm(range(starting_epoch, args.epoch)):
logger.info("=" * 20 + "Training" + "=" * 20)
fixed_cnn.train()
train_loss_meters = AverageMeter()
train_acc_meters = AverageMeter()
train_acc5_meters = AverageMeter()
for images, labels in tqdm(data_loader["train_dataset"]):
images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)
start = time.time()
fixed_cnn.zero_grad()
fixed_cnn_optmizer.zero_grad()
pred = fixed_cnn(images)
loss = criterion(pred, labels)
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(fixed_cnn.parameters(), args.grad_bound)
fixed_cnn_optmizer.step()
acc, acc5 = accuracy(pred, labels, topk=(1, 5))
train_loss_meters.update(loss.item(), labels.shape[0])
train_acc_meters.update(acc.item(), labels.shape[0])
train_acc5_meters.update(acc5.item(), labels.shape[0])
end = time.time()
GLOBAL_STEP += 1
if GLOBAL_STEP % args.train_log_every == 0:
lr = fixed_cnn_optmizer.param_groups[0]['lr']
display = log_display(epoch=epoch,
global_step=GLOBAL_STEP,
time_elapse=end-start,
loss=loss.item(),
loss_avg=train_loss_meters.avg,
acc=acc.item(),
acc_top1_avg=train_acc_meters.avg,
acc_top5_avg=train_acc5_meters.avg,
lr=lr,
gn=grad_norm)
logger.info(display)
if fixed_cnn_scheduler is not None:
fixed_cnn_scheduler.step()
logger.info("="*20 + "Eval" + "="*20)
curr_acc, _ = model_eval(epoch, fixed_cnn, data_loader)
logger.info("curr_acc\t%.4f" % curr_acc)
logger.info("BEST_ACC\t%.4f" % EVAL_BEST_ACC)
payload = '=' * 10 + '\n'
payload = payload + ("curr_acc: %.4f\n best_acc: %.4f\n" % (curr_acc, EVAL_BEST_ACC))
EVAL_BEST_ACC = max(curr_acc, EVAL_BEST_ACC)
TRAIN_HISTORY["train_loss"].append(train_loss_meters.avg)
TRAIN_HISTORY["train_acc"].append(train_acc_meters.avg)
TRAIN_HISTORY["test_acc"].append(curr_acc)
TRAIN_HISTORY["test_acc_best"] = [EVAL_BEST_ACC]
with open(args.checkpoint_path + args.version + '.pickle', 'wb') as handle:
pickle.dump(TRAIN_HISTORY, handle, protocol=pickle.HIGHEST_PROTOCOL)
logger.info("Saved!\n")
return
def model_eval(epoch, fixed_cnn, data_loader, dataset_type='test_dataset'):
global EVAL_STEP
fixed_cnn.eval()
valid_loss_meters = AverageMeter()
valid_acc_meters = AverageMeter()
valid_acc5_meters = AverageMeter()
ce_loss = torch.nn.CrossEntropyLoss()
for images, labels in tqdm(data_loader[dataset_type]):
start = time.time()
images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)
with torch.no_grad():
pred = fixed_cnn(images)
loss = ce_loss(pred, labels)
acc, acc5 = accuracy(pred, labels, topk=(1, 5))
valid_loss_meters.update(loss.item(), labels.shape[0])
valid_acc_meters.update(acc.item(), labels.shape[0])
valid_acc5_meters.update(acc5.item(), labels.shape[0])
end = time.time()
EVAL_STEP += 1
if EVAL_STEP % args.train_log_every == 0:
display = log_display(epoch=epoch,
global_step=GLOBAL_STEP,
time_elapse=end-start,
loss=loss.item(),
test_loss_avg=valid_loss_meters.avg,
acc=acc.item(),
test_acc_avg=valid_acc_meters.avg,
test_acc_top5_avg=valid_acc5_meters.avg)
logger.info(display)
display = log_display(epoch=epoch,
global_step=GLOBAL_STEP,
time_elapse=end-start,
loss=loss.item(),
test_loss_avg=valid_loss_meters.avg,
acc=acc.item(),
test_acc_avg=valid_acc_meters.avg,
test_acc_top5_avg=valid_acc5_meters.avg)
logger.info(display)
return valid_acc_meters.avg, valid_acc5_meters.avg
def do_train(cfg, model, data_loader, loss_factory, optimizer, epoch,
output_dir, tb_log_dir, writer_dict):
logger = logging.getLogger("Training")
batch_time = AverageMeter()
data_time = AverageMeter()
heatmaps_loss_meter = [AverageMeter() for _ in range(cfg.LOSS.NUM_STAGES)]
offset_loss_meter = [AverageMeter() for _ in range(cfg.LOSS.NUM_STAGES)]
model.train()
end = time.time()
for i, (images, heatmaps, masks, offsets,
weights) in enumerate(data_loader):
data_time.update(time.time() - end)
heatmaps = [
list(map(lambda x: x.cuda(non_blocking=True), heatmap))
for heatmap in heatmaps
]
masks = [
list(map(lambda x: x.cuda(non_blocking=True), mask))
for mask in masks
]
offsets = [
list(map(lambda x: x.cuda(non_blocking=True), offset))
for offset in offsets
]
offset_weights = [
list(map(lambda x: x.cuda(non_blocking=True), weight))
for weight in weights
]
####################################################################
if cfg.LOSS.HEATMAP_MIDDLE_LOSS:
outputs, poffsets, middle_output = model(images)
heatmaps_losses, offset_losses, middle_losses = \
loss_factory(outputs, poffsets, heatmaps,
masks, offsets, offset_weights, middle_output)
else:
outputs, poffsets = model(images)
heatmaps_losses, offset_losses = \
loss_factory(outputs, poffsets, heatmaps,
masks, offsets, offset_weights)
####################################################################
loss = 0
for idx in range(cfg.LOSS.NUM_STAGES):
if heatmaps_losses[idx] is not None:
heatmaps_loss = heatmaps_losses[idx].mean(dim=0)
heatmaps_loss_meter[idx].update(heatmaps_loss.item(),
images.size(0))
loss = loss + heatmaps_loss
if offset_losses[idx] is not None:
offset_loss = offset_losses[idx]
offset_loss_meter[idx].update(offset_loss.item(),
images.size(0))
loss = loss + offset_loss
########################################################################
if cfg.LOSS.HEATMAP_MIDDLE_LOSS:
if middle_losses is not None:
loss = loss + middle_losses.mean(dim=0)
########################################################################
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
if i % cfg.PRINT_FREQ == 0 and cfg.RANK == 0:
msg = 'Epoch: [{0}][{1}/{2}]\t' \
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
'Speed: {speed:.1f} samples/s\t' \
'Data: {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
'{heatmaps_loss}{offset_loss}'.format(
epoch, i, len(data_loader),
batch_time=batch_time,
speed=images.size(0)/batch_time.val,
data_time=data_time,
heatmaps_loss=_get_loss_info(
heatmaps_loss_meter, 'heatmaps'),
offset_loss=_get_loss_info(offset_loss_meter, 'offset')
)
logger.info(msg)
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
for idx in range(cfg.LOSS.NUM_STAGES):
writer.add_scalar('train_stage{}_heatmaps_loss'.format(i),
heatmaps_loss_meter[idx].val, global_steps)
writer.add_scalar('train_stage{}_offset_loss'.format(idx),
offset_loss_meter[idx].val, global_steps)
# 每经过PRINT_FREQ个batch,train_global_steps加1
writer_dict['train_global_steps'] = global_steps + 1
prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
for scale_idx in range(len(cfg.DATASET.OUTPUT_SIZE)):
prefix_scale = prefix + '_output_{}'.format(
cfg.DATASET.OUTPUT_SIZE[scale_idx])
save_debug_images(cfg, images, heatmaps[scale_idx],
masks[scale_idx], outputs[scale_idx],
prefix_scale)
def adaptive_quantization(model, val_loader, checkpoint, quantizable_ind,
quantize_bit_choice):
module_list = list(model.modules())
model.load_state_dict(checkpoint) # restore weight first
org_acc = validate(val_loader, model)
# 1. Calculate t_i
# calculate the mean value of adversarial noise for the dataset, notice that we do not add Softmax to network
mean_adv_noise_meter = AverageMeter()
with torch.no_grad():
for input_w, target in val_loader:
input_var = torch.autograd.Variable(input_w).cuda()
output = model(input_var)
top2, inds = torch.topk(output, 2, dim=1)
mean_adv_noise = torch.mean((top2[:, 0] - top2[:, 1])**2 / 2.)
mean_adv_noise_meter.update(mean_adv_noise.cpu().data[0],
output.size(0))
mean_adv_noise_dset = mean_adv_noise_meter.avg
print('Mean adversarial noise for the dataset is: {:.4f}'.format(
mean_adv_noise_dset))
d_acc = 10. # choose 10% for delta_acc. Does not matter
t_i_list = []
for ind in quantizable_ind:
layer = module_list[ind]
assert hasattr(layer, 'weight')
r_W_i_pi = torch.rand(layer.weight.data.size()).cuda(
) - 0.5 # re-normalize to [-0.5, 0.5]
k = k_min = 1e-5
k_max = 1e1
# get initial acc'
model.load_state_dict(checkpoint) # restore weight first
# assert validate(val_loader, model) == org_acc # removed to accelerate
layer.weight.data += k * r_W_i_pi
new_acc = validate(val_loader, model)
while not np.abs(org_acc - new_acc - d_acc) < 0.1:
if org_acc - new_acc < d_acc:
k_min = k
else:
k_max = k
k = np.sqrt(k_min * k_max).item()
model.load_state_dict(checkpoint) # restore weight first
layer.weight.data += k * r_W_i_pi
new_acc = validate(val_loader, model)
print('Layer {} current acc degradation: {:.3f}'.format(
ind, org_acc - new_acc))
mean_r_z_i = AverageMeter()
with torch.no_grad():
for i, (input_d, target) in enumerate(val_loader):
input_var = torch.autograd.Variable(input_d).cuda()
# compute output
model.load_state_dict(checkpoint) # restore weight first
output1 = model(input_var)
layer.weight.data += k * r_W_i_pi
output2 = model(input_var)
norm_r = torch.norm(output1 - output2, p=2, dim=1)**2
mean_r_z_i.update(
torch.mean(norm_r).cpu().data[0], output1.size(0))
t_i = mean_r_z_i.avg / mean_adv_noise_dset
print('==> t_i for layer {}: {}'.format(ind, t_i))
t_i_list.append(t_i)
print('t_i_list: ')
print(t_i_list)
# t_i_list = [113.314645623, 108.02437323, 111.228006385, 109.585273767, 115.362011096, 111.136186881, 114.150737539,
# 106.789374298, 135.436417323, 118.175965146, 136.776404035, 162.089406771, 224.905060191, 340.589419784,
# 904.878690392, 256.250864841]
# 2. Calculate p_i
fix_b_i = 6
p_i_list = []
for i, ind in enumerate(quantizable_ind):
model.load_state_dict(checkpoint) # restore weight first
centroids, labels = k_means_torch_compact(module_list[ind].weight.data,
2**fix_b_i,
init='linear',
use_gpu=False)
w_q = reconstruct_weight_from_k_means_result(centroids, labels)
del centroids, labels
mean_r_z_i = AverageMeter()
with torch.no_grad():
for input_d, target in val_loader:
input_var = torch.autograd.Variable(input_d).cuda()
# compute output
model.load_state_dict(checkpoint) # restore weight first
output1 = model(input_var)
del module_list[ind].weight
module_list[ind].weight = nn.Parameter(w_q.float())
output2 = model(input_var)
norm_r = torch.norm(output1 - output2, p=2, dim=1)**2
mean_r_z_i.update(
torch.mean(norm_r).cpu().data[0], output1.size(0))
del w_q
p_i = mean_r_z_i.avg # / np.exp(-np.log(4) * fix_b_i)
print('==> p_i for layer {}: {}'.format(ind, p_i))
p_i_list.append(p_i)
# 3. Calculate b_i
b1 = 8
assert len(p_i_list) == len(t_i_list) == len(quantizable_ind)
layer1_size = np.prod(module_list[quantizable_ind[0]].weight.data.size())
bits_list = [b1]
print('==> Layer1 size: {}, bits: {}'.format(layer1_size, b1))
for i, ind in enumerate(quantizable_ind):
if i == 0:
continue
this_size = np.prod(module_list[ind].weight.data.size())
b_i = b1 + (1 / np.log(4)) * np.log(
p_i_list[i] * t_i_list[0] * layer1_size /
(p_i_list[0] * t_i_list[i] * this_size))
print('Optimal bits for layer {}: {}'.format(ind, b_i))
bits_list.append(b_i)
print('Final result: {}'.format(bits_list))
def train(**kwargs):
opt.parse(kwargs)
if not os.path.exists(opt.save_folder):
os.mkdir(opt.save_folder)
tb_logger = SummaryWriter(opt.save_folder)
logger = create_logger('global_logger', opt.save_folder + '/log.txt')
batch_time = AverageMeter(10)
data_time = AverageMeter(10)
losses = AverageMeter(10)
loss_meter = meter.AverageValueMeter()
train_sets = []
for data_txt in opt.train_txt:
data_root, gt_root, list_file = data_txt.split(' ')
train_sets.append(
OCRDataset(data_root, gt_root, list_file, opt.input_size, 'train',
opt.chars_list, opt.max_seq))
train_data = ConcatDataset(train_sets)
train_loader = DataLoader(train_data,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_works)
valid_sets = []
for valid_txt in opt.valid_txt:
data_root, gt_root, list_file = valid_txt.split(' ')
valid_sets.append(
OCRDataset(data_root, gt_root, list_file, opt.input_size, 'valid',
opt.chars_list, opt.max_seq))
valid_data = ConcatDataset(valid_sets)
valid_loader = DataLoader(valid_data,
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.num_works)
model = getattr(models, opt.model)(opt.basenet,
opt.input_size,
opt.max_seq,
opt.num_classes,
mode='train',
attn=opt.attn)
if opt.load_model_path is not None:
load_state(model, opt.load_model_path, 'cuda:%d' % opt.gpus[0])
if len(opt.gpus) > 1:
model = torch.nn.DataParallel(model, device_ids=opt.gpus)
model = gpu(model, opt)
if len(opt.gpus) > 1:
optimizer = torch.optim.Adam(model.module.parameters(),
lr=opt.lr,
betas=opt.betas,
weight_decay=opt.weight_decay)
else:
optimizer = torch.optim.Adam(model.parameters(),
lr=opt.lr,
betas=opt.betas,
weight_decay=opt.weight_decay)
curr_step = 0
total_step = int(len(train_data) / opt.batch_size * opt.epoches)
best_val_error = 1e10
previous_loss = 1e10
# warmup
warmup_epoches = opt.epoches // 10
warmup_rate = math.pow(100, 1 / warmup_epoches)
for epoch in range(opt.epoches):
model.train()
end = time.time()
# loss_meter.reset()
for i, (imgs, gt_chars_seg, gt_order_seg,
gt_pos_seg) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# zero the parameter gradients
optimizer.zero_grad()
imgs = gpu(imgs, opt)
gt_chars_seg = gpu(gt_chars_seg, opt)
gt_order_seg = gpu(gt_order_seg, opt)
gt_pos_seg = gpu(gt_pos_seg, opt)
chars_seg, ord_seg, pos_seg = model(imgs)
loss = get_loss(chars_seg, ord_seg, pos_seg, gt_chars_seg,
gt_order_seg, gt_pos_seg, opt)
loss.backward()
optimizer.step()
losses.update(loss.item())
loss_meter.add(loss.item())
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
curr_step += 1
current_lr = optimizer.param_groups[0]['lr']
if curr_step % opt.print_freq == 0:
tb_logger.add_scalar('loss_train', losses.avg, curr_step)
tb_logger.add_scalar('lr', current_lr, curr_step)
logger.info(
'Iter: [{0}/{1}]\t'
'Epoch: {2}\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'LR {lr:.4f}'.format(curr_step,
total_step,
epoch,
batch_time=batch_time,
data_time=data_time,
loss=losses,
lr=current_lr))
# val
model.eval()
val_error = val(model, valid_loader, opt)
logger.info('Mean error: {0}\t'.format(val_error))
if not tb_logger is None:
tb_logger.add_scalar('error_val', val_error, curr_step)
if val_error < best_val_error:
best_val_error = val_error
if len(opt.gpus) > 1:
torch.save(model.module.state_dict(),
os.path.join(opt.save_folder, "best_val_error.pth"))
else:
torch.save(model.state_dict(),
os.path.join(opt.save_folder, "best_val_error.pth"))
# warmup
if epoch < warmup_epoches:
for param_group in optimizer.param_groups:
param_group["lr"] *= warmup_rate
# decay lr if loss no longer decrease
else:
if opt.lr_immediate_decay and loss_meter.value(
)[0] > previous_loss:
for param_group in optimizer.param_groups:
param_group["lr"] *= opt.lr_decay
if epoch == int(opt.epoches * 0.6) or epoch == int(
opt.epoches * 0.9):
for param_group in optimizer.param_groups:
param_group["lr"] *= opt.lr_decay
previous_loss = loss_meter.value()[0]
# save last pth
if len(opt.gpus) > 1:
torch.save(model.module.state_dict(),
os.path.join(opt.save_folder, "last.pth"))
else:
torch.save(model.state_dict(), os.path.join(opt.save_folder,
"last.pth"))
def validate(val_loader, model):
from utils import accuracy
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
criterion = nn.CrossEntropyLoss().cuda()
# switch to evaluate mode
model.eval()
with torch.no_grad():
for i, (input, target) in enumerate(val_loader):
target = target.cuda(async=True)
input_var = torch.autograd.Variable(input).cuda()
target_var = torch.autograd.Variable(target).cuda()
# compute output
output = model(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
losses.update(loss.data[0], input.size(0))
top1.update(prec1[0], input.size(0))
top5.update(prec5[0], input.size(0))
return top1.avg
def validate(val_loader, miml, encoder, decoder, criterion, epoch, writer):
"""
Performs one epoch's validation.
:param val_loader: DataLoader for validation data.
:param encoder: encoder model
:param decoder: decoder model
:param criterion: loss layer
:return: BLEU-4 score
"""
miml.eval()
encoder.eval()
decoder.eval() # eval mode (no dropout or batchnorm)
total_step = len(val_loader)
batch_time = AverageMeter()
losses = AverageMeter()
top5accs = AverageMeter()
start = time.time()
references = list(
) # references (true captions) for calculating BLEU-4 score
hypotheses = list() # hypotheses (predictions)
with torch.no_grad():
# Batches
for i, (imgs, caps, caplens, allcaps) in enumerate(val_loader):
# Move to device, if available
imgs = imgs.to(device)
caps = caps.to(device)
caplens = caplens.to(device)
# Forward prop.
attrs = miml(imgs)
imgs = encoder(imgs)
scores, caps_sorted, decode_lengths, alphas, sort_ind = decoder(
attrs, imgs, caps, caplens)
# Since we decoded starting with <start>, the targets are all words after <start>, up to <end>
targets = caps_sorted[:, 1:]
# Remove timesteps that we didn't decode at, or are pads
# pack_padded_sequence is an easy trick to do this
scores_copy = scores.clone()
scores, _ = pack_padded_sequence(scores,
decode_lengths,
batch_first=True)
targets, _ = pack_padded_sequence(targets,
decode_lengths,
batch_first=True)
# Calculate loss
loss = criterion(scores, targets)
loss += alpha_c * ((1. - alphas.sum(dim=1))**2).mean()
# Keep track of metrics
losses.update(loss.item(), sum(decode_lengths))
top5 = accuracy(scores, targets, 5)
top5accs.update(top5, sum(decode_lengths))
batch_time.update(time.time() - start)
start = time.time()
if i % print_freq == 0:
writer.add_scalars('val', {
'loss': loss.item(),
'mAp': top5accs.val
}, epoch * total_step + i)
print(
'Validation: [{0}/{1}]\t'
'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Top-5 Accuracy {top5.val:.3f} ({top5.avg:.3f})\t'.format(
i,
len(val_loader),
batch_time=batch_time,
loss=losses,
top5=top5accs))
# Store references (true captions), and hypothesis (prediction) for each image
# If for n images, we have n hypotheses, and references a, b, c... for each image, we need -
# references = [[ref1a, ref1b, ref1c], [ref2a, ref2b], ...], hypotheses = [hyp1, hyp2, ...]
# References
# because images were sorted in the decoder
allcaps = allcaps[sort_ind]
for j in range(allcaps.shape[0]):
img_caps = allcaps[j].tolist()
img_captions = list(
map(
lambda c: [
w for w in c if w not in
{word_map['<start>'], word_map['<pad>']}
], img_caps)) # remove <start> and pads
references.append(img_captions)
# Hypotheses
_, preds = torch.max(scores_copy, dim=2)
preds = preds.tolist()
temp_preds = list()
for j, p in enumerate(preds):
temp_preds.append(preds[j][:decode_lengths[j]]) # remove pads
preds = temp_preds
hypotheses.extend(preds)
assert len(references) == len(hypotheses)
# Calculate BLEU-4 scores
weights = (1.0 / 1.0, )
bleu1 = corpus_bleu(references, hypotheses, weights)
weights = (
1.0 / 2.0,
1.0 / 2.0,
)
bleu2 = corpus_bleu(references, hypotheses, weights)
weights = (
1.0 / 3.0,
1.0 / 3.0,
1.0 / 3.0,
)
bleu3 = corpus_bleu(references, hypotheses, weights)
bleu4 = corpus_bleu(references, hypotheses)
writer.add_scalars('Bleu', {
'Bleu1': bleu1,
'Bleu2': bleu2,
'Bleu3': bleu3,
'Bleu4': bleu4
}, epoch)
print(
'\n * LOSS - {loss.avg:.3f}, TOP-5 ACCURACY - {top5.avg:.3f}, BLEU-4 - {bleu}\n'
.format(loss=losses, top5=top5accs, bleu=bleu4))
return bleu4
def test(model, test_dataset, test_loader, output_filename, result_dir=None):
load_batch_size = test_loader.batch_size
num_batches = len(test_loader)
model.eval()
fw_time_meter = AverageMeter()
det_results = {}
gt_results = {}
for i, (data_dicts, datas) in enumerate(test_loader):
torch.cuda.synchronize()
tic = time.time()
if data_dicts is None:
continue
predictions = predict(model, data_dicts)
torch.cuda.synchronize()
fw_time_meter.update((time.time() - tic))
print('%d/%d %.3f' % (i, num_batches, fw_time_meter.val))
# datas = test_dataset.get_frustum_data(i)
# print(f'Datas len: {len(datas)}')
# print(f'Predictions len: {len(predictions)}')
for data, pred in zip(datas, predictions):
if data is None:
continue
data_idx = data['idx_i']
class_type = data['type_i']
box2d = data['box2d_i']
box3d = data['box3d_i']
box3d_sizes = data['box3d_size_i']
ry_gt = data['heading_i']
box_pos_gt = data['obj_pos_i']
if data_idx not in det_results:
det_results[data_idx] = {}
gt_results[data_idx] = {}
if class_type not in det_results[data_idx]:
det_results[data_idx][class_type] = []
gt_results[data_idx][class_type] = []
x1, y1, x2, y2 = box2d
l_gt, w_gt, h_gt = box3d_sizes
tx_gt, ty_gt, tz_gt = box_pos_gt
output_gt = [
x1, y1, x2, y2, tx_gt, ty_gt, tz_gt, h_gt, w_gt, l_gt, ry_gt
]
gt_results[data_idx][class_type].append(output_gt)
# print('****************')
# print(tx_gt, ty_gt, tz_gt, h_gt, w_gt, l_gt, ry_gt)
# print('================')
for n in range(len(pred)):
h, w, l, tx, ty, tz, ry, score = pred[n]
output = [x1, y1, x2, y2, tx, ty, tz, h, w, l, ry, score]
# print(tx, ty, tz, h, w, l, ry)
# output = [x1, y1, x2, y2, tx, ty, tz, h, w, l, ry,1]
# print(score)
det_results[data_idx][class_type].append(output)
# print('****************')
num_images = len(det_results)
logging.info('Average time:')
logging.info('batch:%0.3f' % fw_time_meter.avg)
logging.info('avg_per_object:%0.3f' %
(fw_time_meter.avg / load_batch_size))
logging.info('avg_per_image:%.3f' %
(fw_time_meter.avg * len(test_loader) / num_images))
return gt_results, det_results
def train(args, train_loader, model, optimizer, epoch, curr_lr, win_feats5,
win_fusion, viz, global_step, accumulation_steps):
batch_time = AverageMeter()
data_time = AverageMeter()
feats5_losses = AverageMeter()
fusion_losses = AverageMeter()
total_losses = AverageMeter()
# switch to eval mode to make BN unchanged.
model.train()
optimizer.zero_grad()
end = time.time()
for i, (img, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# Input for Image CNN.
img_var = utils.check_gpu(0, img) # BS X 3 X H X W
target_var = utils.check_gpu(0, target) # BS X H X W X NUM_CLASSES
bs = img.size()[0] * accumulation_steps
score_feats5, fused_feats = model(
img_var) # BS X NUM_CLASSES X 472 X 472
feats5_loss = WeightedMultiLabelSigmoidLoss(score_feats5, target_var)
fused_feats_loss = WeightedMultiLabelSigmoidLoss(
fused_feats, target_var)
loss = feats5_loss + fused_feats_loss
loss.backward()
# clear memory
del img_var
del target_var
del score_feats5
torch.cuda.empty_cache()
# increase batch size by factor of accumulation steps (Gradient accumulation) for training with limited memory
if (i + 1) % accumulation_steps == 0:
feats5_losses.update(feats5_loss.data, bs)
fusion_losses.update(fused_feats_loss.data, bs)
total_losses.update(loss.data, bs)
# Only plot the fused feats loss.
trn_feats5_loss = feats5_loss.clone().cpu().data.numpy()
trn_fusion_loss = fused_feats_loss.clone().cpu().data.numpy()
viz.line(win=win_feats5,
name='train_feats5',
update='append',
X=np.array([global_step]),
Y=np.array([trn_feats5_loss]))
viz.line(win=win_fusion,
name='train_fusion',
update='append',
X=np.array([global_step]),
Y=np.array([trn_fusion_loss]))
optimizer.step()
optimizer.zero_grad()
global_step += 1
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if ((i + 1) % args.print_freq == 0):
print("\n\n")
print(
'Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Total Loss {total_loss.val:.11f} ({total_loss.avg:.11f})\n'
'lr {learning_rate:.10f}\t'.format(
epoch,
int((i + 1) / accumulation_steps),
int(len(train_loader) / accumulation_steps),
batch_time=batch_time,
data_time=data_time,
total_loss=total_losses,
learning_rate=curr_lr))
del feats5_loss
del fused_feats_loss
del feats5_losses
del fusion_losses
del total_losses
torch.cuda.empty_cache()
return global_step
def train_net(param, model, train_data, valid_data, plot=False,device='cuda'):
# 初始化参数
model_name = param['model_name']
epochs = param['epochs']
batch_size = param['batch_size']
lr = param['lr']
gamma = param['gamma']
step_size = param['step_size']
momentum = param['momentum']
weight_decay = param['weight_decay']
disp_inter = param['disp_inter']
save_inter = param['save_inter']
min_inter = param['min_inter']
iter_inter = param['iter_inter']
save_log_dir = param['save_log_dir']
save_ckpt_dir = param['save_ckpt_dir']
load_ckpt_dir = param['load_ckpt_dir']
#
scaler = GradScaler()
# 网络参数
train_data_size = train_data.__len__()
valid_data_size = valid_data.__len__()
c, y, x = train_data.__getitem__(0)['image'].shape
train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True, num_workers=1)
valid_loader = DataLoader(dataset=valid_data, batch_size=batch_size, shuffle=False, num_workers=1)
optimizer = optim.AdamW(model.parameters(), lr=3e-4 ,weight_decay=weight_decay)
#optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=momentum, weight_decay=weight_decay)
#scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=3, T_mult=2, eta_min=1e-5, last_epoch=-1)
#criterion = nn.CrossEntropyLoss(reduction='mean').to(device)
DiceLoss_fn=DiceLoss(mode='multiclass')
SoftCrossEntropy_fn=SoftCrossEntropyLoss(smooth_factor=0.1)
criterion = L.JointLoss(first=DiceLoss_fn, second=SoftCrossEntropy_fn,
first_weight=0.5, second_weight=0.5).cuda()
logger = inial_logger(os.path.join(save_log_dir, time.strftime("%m-%d %H:%M:%S", time.localtime()) +'_'+model_name+ '.log'))
# 主循环
train_loss_total_epochs, valid_loss_total_epochs, epoch_lr = [], [], []
train_loader_size = train_loader.__len__()
valid_loader_size = valid_loader.__len__()
best_iou = 0
best_epoch=0
best_mode = copy.deepcopy(model)
epoch_start = 0
if load_ckpt_dir is not None:
ckpt = torch.load(load_ckpt_dir)
epoch_start = ckpt['epoch']
model.load_state_dict(ckpt['state_dict'])
optimizer.load_state_dict(ckpt['optimizer'])
logger.info('Total Epoch:{} Image_size:({}, {}) Training num:{} Validation num:{}'.format(epochs, x, y, train_data_size, valid_data_size))
#
for epoch in range(epoch_start, epochs):
epoch_start = time.time()
# 训练阶段
model.train()
train_epoch_loss = AverageMeter()
train_iter_loss = AverageMeter()
for batch_idx, batch_samples in enumerate(train_loader):
data, target = batch_samples['image'], batch_samples['label']
data, target = Variable(data.to(device)), Variable(target.to(device))
with autocast(): #need pytorch>1.6
pred = model(data)
loss = criterion(pred, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
scheduler.step(epoch + batch_idx / train_loader_size)
image_loss = loss.item()
train_epoch_loss.update(image_loss)
train_iter_loss.update(image_loss)
if batch_idx % iter_inter == 0:
spend_time = time.time() - epoch_start
logger.info('[train] epoch:{} iter:{}/{} {:.2f}% lr:{:.6f} loss:{:.6f} ETA:{}min'.format(
epoch, batch_idx, train_loader_size, batch_idx/train_loader_size*100,
optimizer.param_groups[-1]['lr'],
train_iter_loss.avg,spend_time / (batch_idx+1) * train_loader_size // 60 - spend_time // 60))
train_iter_loss.reset()
# 验证阶段
model.eval()
valid_epoch_loss = AverageMeter()
valid_iter_loss = AverageMeter()
iou=IOUMetric(10)
with torch.no_grad():
for batch_idx, batch_samples in enumerate(valid_loader):
data, target = batch_samples['image'], batch_samples['label']
data, target = Variable(data.to(device)), Variable(target.to(device))
pred = model(data)
loss = criterion(pred, target)
pred=pred.cpu().data.numpy()
pred= np.argmax(pred,axis=1)
iou.add_batch(pred,target.cpu().data.numpy())
#
image_loss = loss.item()
valid_epoch_loss.update(image_loss)
valid_iter_loss.update(image_loss)
# if batch_idx % iter_inter == 0:
# logger.info('[val] epoch:{} iter:{}/{} {:.2f}% loss:{:.6f}'.format(
# epoch, batch_idx, valid_loader_size, batch_idx / valid_loader_size * 100, valid_iter_loss.avg))
val_loss=valid_iter_loss.avg
acc, acc_cls, iu, mean_iu, fwavacc=iou.evaluate()
logger.info('[val] epoch:{} miou:{:.2f}'.format(epoch,mean_iu))
# 保存loss、lr
train_loss_total_epochs.append(train_epoch_loss.avg)
valid_loss_total_epochs.append(valid_epoch_loss.avg)
epoch_lr.append(optimizer.param_groups[0]['lr'])
# 保存模型
if epoch % save_inter == 0 and epoch > min_inter:
state = {'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
filename = os.path.join(save_ckpt_dir, 'checkpoint-epoch{}.pth'.format(epoch))
torch.save(state, filename) # pytorch1.6会压缩模型，低版本无法加载
# 保存最优模型
if mean_iu > best_iou: # train_loss_per_epoch valid_loss_per_epoch
state = {'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
filename = os.path.join(save_ckpt_dir, 'checkpoint-best.pth')
torch.save(state, filename)
best_iou = mean_iu
best_mode = copy.deepcopy(model)
logger.info('[save] Best Model saved at epoch:{} ============================='.format(epoch))
#scheduler.step()
# 显示loss
# 训练loss曲线
if plot:
x = [i for i in range(epochs)]
fig = plt.figure(figsize=(12, 4))
ax = fig.add_subplot(1, 2, 1)
ax.plot(x, smooth(train_loss_total_epochs, 0.6), label='train loss')
ax.plot(x, smooth(valid_loss_total_epochs, 0.6), label='val loss')
ax.set_xlabel('Epoch', fontsize=15)
ax.set_ylabel('CrossEntropy', fontsize=15)
ax.set_title('train curve', fontsize=15)
ax.grid(True)
plt.legend(loc='upper right', fontsize=15)
ax = fig.add_subplot(1, 2, 2)
ax.plot(x, epoch_lr, label='Learning Rate')
ax.set_xlabel('Epoch', fontsize=15)
ax.set_ylabel('Learning Rate', fontsize=15)
ax.set_title('lr curve', fontsize=15)
ax.grid(True)
plt.legend(loc='upper right', fontsize=15)
plt.show()
return best_mode, model
def test(model, test_dataset, test_loader, output_filename, result_dir=None):
load_batch_size = test_loader.batch_size
num_batches = len(test_loader)
model.eval()
fw_time_meter = AverageMeter()
det_results = {}
for i, data_dicts in enumerate(test_loader):
point_clouds = data_dicts['point_cloud']
rot_angles = data_dicts['rot_angle']
# optional
ref_centers = data_dicts.get('ref_center')
rgb_probs = data_dicts.get('rgb_prob')
# from ground truth box detection
if rgb_probs is None:
rgb_probs = torch.ones_like(rot_angles)
# not belong to refinement stage
if ref_centers is None:
ref_centers = torch.zeros((point_clouds.shape[0], 3))
batch_size = point_clouds.shape[0]
rot_angles = rot_angles.view(-1)
rgb_probs = rgb_probs.view(-1)
if 'box3d_center' in data_dicts:
data_dicts.pop('box3d_center')
data_dicts_var = {
key: value.cuda()
for key, value in data_dicts.items()
}
torch.cuda.synchronize()
tic = time.time()
with torch.no_grad():
outputs = model(data_dicts_var)
cls_probs, center_preds, heading_preds, size_preds = outputs
torch.cuda.synchronize()
fw_time_meter.update((time.time() - tic))
num_pred = cls_probs.shape[1]
print('%d/%d %.3f' % (i, num_batches, fw_time_meter.val))
cls_probs = cls_probs.data.cpu().numpy()
center_preds = center_preds.data.cpu().numpy()
heading_preds = heading_preds.data.cpu().numpy()
size_preds = size_preds.data.cpu().numpy()
rgb_probs = rgb_probs.numpy()
rot_angles = rot_angles.numpy()
ref_centers = ref_centers.numpy()
for b in range(batch_size):
if cfg.TEST.METHOD == 'nms':
fg_idx = (cls_probs[b, :, 0] < cls_probs[b, :, 1]).nonzero()[0]
if fg_idx.size == 0:
fg_idx = np.argmax(cls_probs[b, :, 1])
fg_idx = np.array([fg_idx])
else:
fg_idx = np.argmax(cls_probs[b, :, 1])
fg_idx = np.array([fg_idx])
num_pred = len(fg_idx)
single_centers = center_preds[b, fg_idx]
single_headings = heading_preds[b, fg_idx]
single_sizes = size_preds[b, fg_idx]
single_scores = cls_probs[b, fg_idx, 1] + rgb_probs[b]
data_idx = test_dataset.id_list[load_batch_size * i + b]
class_type = test_dataset.type_list[load_batch_size * i + b]
box2d = test_dataset.box2d_list[load_batch_size * i + b]
rot_angle = rot_angles[b]
ref_center = ref_centers[b]
if data_idx not in det_results:
det_results[data_idx] = {}
if class_type not in det_results[data_idx]:
det_results[data_idx][class_type] = []
for n in range(num_pred):
x1, y1, x2, y2 = box2d
score = single_scores[n]
h, w, l, tx, ty, tz, ry = from_prediction_to_label_format(
single_centers[n], single_headings[n], single_sizes[n],
rot_angle, ref_center)
output = [x1, y1, x2, y2, tx, ty, tz, h, w, l, ry, score]
det_results[data_idx][class_type].append(output)
num_images = len(det_results)
logging.info('Average time:')
logging.info('batch:%0.3f' % fw_time_meter.avg)
logging.info('avg_per_object:%0.3f' %
(fw_time_meter.avg / load_batch_size))
logging.info('avg_per_image:%.3f' %
(fw_time_meter.avg * len(test_loader) / num_images))
# Write detection results for KITTI evaluation
if cfg.TEST.METHOD == 'nms':
write_detection_results_nms(result_dir, det_results)
else:
write_detection_results(result_dir, det_results)
output_dir = os.path.join(result_dir, 'data')
if 'test' not in cfg.TEST.DATASET:
evaluate_py_wrapper(result_dir)
# evaluate_cuda_wrapper(output_dir, cfg.TEST.DATASET)
else:
logger.info('results file save in {}'.format(result_dir))
os.system('cd %s && zip -q -r ../results.zip *' % (result_dir))
def train_val(model, args):
train_dir = args.train_dir
val_dir = args.val_dir
config = Config(args.config)
cudnn.benchmark = True
# train
train_loader = torch.utils.data.DataLoader(lsp_lspet_data.LSP_Data(
'lspet', train_dir, 8,
Mytransforms.Compose([
Mytransforms.RandomResized(),
Mytransforms.RandomRotate(40),
Mytransforms.RandomCrop(368),
Mytransforms.RandomHorizontalFlip(),
])),
batch_size=config.batch_size,
shuffle=True,
num_workers=config.workers,
pin_memory=True)
# val
if args.val_dir is not None and config.test_interval != 0:
# val
val_loader = torch.utils.data.DataLoader(lsp_lspet_data.LSP_Data(
'lsp', val_dir, 8,
Mytransforms.Compose([
Mytransforms.TestResized(368),
])),
batch_size=config.batch_size,
shuffle=True,
num_workers=config.workers,
pin_memory=True)
if args.gpu[0] < 0:
criterion = nn.MSELoss()
else:
criterion = nn.MSELoss().cuda()
params, multiple = get_parameters(model, config, True)
# params, multiple = get_parameters(model, config, False)
optimizer = torch.optim.SGD(params,
config.base_lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
losses_list = [AverageMeter() for i in range(6)]
end = time.time()
iters = config.start_iters
best_model = config.best_model
heat_weight = 46 * 46 * 15 / 1.0
while iters < config.max_iter:
for i, (input, heatmap, centermap) in enumerate(train_loader):
learning_rate = adjust_learning_rate(
optimizer,
iters,
config.base_lr,
policy=config.lr_policy,
policy_parameter=config.policy_parameter,
multiple=multiple)
data_time.update(time.time() - end)
if args.gpu[0] >= 0:
heatmap = heatmap.cuda(async=True)
centermap = centermap.cuda(async=True)
input_var = torch.autograd.Variable(input)
heatmap_var = torch.autograd.Variable(heatmap)
centermap_var = torch.autograd.Variable(centermap)
heat1, heat2, heat3, heat4, heat5, heat6 = model(
input_var, centermap_var)
loss1 = criterion(heat1, heatmap_var) * heat_weight
loss2 = criterion(heat2, heatmap_var) * heat_weight
loss3 = criterion(heat3, heatmap_var) * heat_weight
loss4 = criterion(heat4, heatmap_var) * heat_weight
loss5 = criterion(heat5, heatmap_var) * heat_weight
loss6 = criterion(heat6, heatmap_var) * heat_weight
loss = loss1 + loss2 + loss3 + loss4 + loss5 + loss6
#print(input.size(0).item())
losses.update(loss.item(), input.size(0))
for cnt, l in enumerate([loss1, loss2, loss3, loss4, loss5,
loss6]):
losses_list[cnt].update(l.item(), input.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
iters += 1
if iters % config.display == 0:
print(
'Train Iteration: {0}\t'
'Time {batch_time.sum:.3f}s / {1}iters, ({batch_time.avg:.3f})\t'
'Data load {data_time.sum:.3f}s / {1}iters, ({data_time.avg:3f})\n'
'Learning rate = {2}\n'
'Loss = {loss.val:.8f} (ave = {loss.avg:.8f})\n'.format(
iters,
config.display,
learning_rate,
batch_time=batch_time,
data_time=data_time,
loss=losses))
for cnt in range(0, 6):
print(
'Loss{0} = {loss1.val:.8f} (ave = {loss1.avg:.8f})\t'.
format(cnt + 1, loss1=losses_list[cnt]))
print(
time.strftime(
'%Y-%m-%d %H:%M:%S -----------------------------------------------------------------------------------------------------------------\n',
time.localtime()))
batch_time.reset()
data_time.reset()
losses.reset()
for cnt in range(6):
losses_list[cnt].reset()
save_checkpoint({
'iter': iters,
'state_dict': model.state_dict(),
}, 0, args.model_name)
# val
if args.val_dir is not None and config.test_interval != 0 and iters % config.test_interval == 0:
model.eval()
for j, (input, heatmap, centermap) in enumerate(val_loader):
if args.cuda[0] >= 0:
heatmap = heatmap.cuda(async=True)
centermap = centermap.cuda(async=True)
input_var = torch.autograd.Variable(input)
heatmap_var = torch.autograd.Variable(heatmap)
centermap_var = torch.autograd.Variable(centermap)
heat1, heat2, heat3, heat4, heat5, heat6 = model(
input_var, centermap_var)
loss1 = criterion(heat1, heatmap_var) * heat_weight
loss2 = criterion(heat2, heatmap_var) * heat_weight
loss3 = criterion(heat3, heatmap_var) * heat_weight
loss4 = criterion(heat4, heatmap_var) * heat_weight
loss5 = criterion(heat5, heatmap_var) * heat_weight
loss6 = criterion(heat6, heatmap_var) * heat_weight
loss = loss1 + loss2 + loss3 + loss4 + loss5 + loss6
losses.update(loss.data[0], input.size(0))
for cnt, l in enumerate(
[loss1, loss2, loss3, loss4, loss5, loss6]):
losses_list[cnt].update(l.data[0], input.size(0))
batch_time.update(time.time() - end)
end = time.time()
is_best = losses.avg < best_model
best_model = min(best_model, losses.avg)
save_checkpoint(
{
'iter': iters,
'state_dict': model.state_dict(),
}, is_best, args.model_name)
if j % config.display == 0:
print(
'Test Iteration: {0}\t'
'Time {batch_time.sum:.3f}s / {1}iters, ({batch_time.avg:.3f})\t'
'Data load {data_time.sum:.3f}s / {1}iters, ({data_time.avg:3f})\n'
'Loss = {loss.val:.8f} (ave = {loss.avg:.8f})\n'.
format(j,
config.display,
batch_time=batch_time,
data_time=data_time,
loss=losses))
for cnt in range(0, 6):
print(
'Loss{0} = {loss1.val:.8f} (ave = {loss1.avg:.8f})\t'
.format(cnt + 1, loss1=losses_list[cnt]))
print(
time.strftime(
'%Y-%m-%d %H:%M:%S -----------------------------------------------------------------------------------------------------------------\n',
time.localtime()))
batch_time.reset()
losses.reset()
for cnt in range(6):
losses_list[cnt].reset()
model.train()
def train(train_loader, net, criterion, optimizer, scheduler, args,
val_loader):
bestaccT = 0
bestfwiou = 0.5
bestaccV = 0.0
bestloss = 1
begin_time = time.time()
all_iters = float(len(train_loader) * args['epochs'])
curr_epoch = 0
while True:
if args['gpu']: torch.cuda.empty_cache()
net.train()
start = time.time()
acc_meter = AverageMeter()
train_main_loss = AverageMeter()
train_aux_loss = AverageMeter()
curr_iter = curr_epoch * len(train_loader)
for i, (imgs, labels) in enumerate(train_loader):
running_iter = curr_iter + i + 1
adjust_lr(optimizer, running_iter, all_iters)
#imgs = torch.squeeze(imgs)
imgs = imgs.float()
labels = labels.long()
#imgs, labels = data
if args['gpu']:
imgs = imgs.cuda().float()
labels = labels.cuda().long()
optimizer.zero_grad()
outputs, aux = net(imgs) #
assert outputs.shape[1] == RS.num_classes + 1
loss_main = criterion(outputs, labels)
loss_aux = criterion(aux, labels)
loss = loss_main * 0.7 + loss_aux * 0.3
loss.backward()
optimizer.step()
labels = labels.cpu().detach().numpy()
outputs = outputs.cpu().detach()
_, preds = torch.max(outputs, dim=1)
preds = preds.numpy()
# batch_valid_sum = 0
acc_curr_meter = AverageMeter()
for (pred, label) in zip(preds, labels):
acc, _ = accuracy(pred, label)
acc_curr_meter.update(acc)
acc_meter.update(acc_curr_meter.avg)
train_main_loss.update(loss.cpu().detach().numpy())
train_aux_loss.update(loss_aux.cpu().detach().numpy())
curr_time = time.time() - start
if (i + 1) % args['print_freq'] == 0:
print(
'[epoch %d] [iter %d / %d %.1fs] [lr %f] [train loss %.4f acc %.2f]'
% (curr_epoch, i + 1, len(train_loader), curr_time,
optimizer.param_groups[0]['lr'], train_main_loss.val,
acc_meter.val * 100))
writer.add_scalar('train_main_loss', train_main_loss.val,
running_iter)
writer.add_scalar('train_accuracy', acc_meter.val,
running_iter)
writer.add_scalar('train_aux_loss', train_aux_loss.avg,
running_iter)
writer.add_scalar('lr', optimizer.param_groups[0]['lr'],
running_iter)
acc_v, fwiou_v, loss_v = validate(val_loader, net, criterion,
curr_epoch)
if acc_meter.avg > bestaccT: bestaccT = acc_meter.avg
if fwiou_v > bestfwiou:
bestfwiou = fwiou_v
bestloss = loss_v
bestaccV = acc_v
torch.save(
net.state_dict(),
os.path.join(
args['chkpt_dir'], NET_NAME + '_%de_OA%.2f_fwIoU%.2f.pth' %
(curr_epoch, acc_v * 100, fwiou_v * 100)))
print(
'Total time: %.1fs Best rec: Train acc %.2f, Val acc %.2f fwiou %.2f, Val_loss %.4f'
% (time.time() - begin_time, bestaccT * 100, bestaccV * 100,
bestfwiou * 100, bestloss))
curr_epoch += 1
#scheduler.step()
if curr_epoch >= args['epochs']:
return
def train_moco(epoch, train_loader, model, model_ema, contrast, criterion,
optimizer, opt, recorder):
"""
one epoch training for instance discrimination
"""
print("==> (MoCo) training...")
model_ema.eval()
model.train()
loss_meter = AverageMeter()
prob_meter = AverageMeter()
batch_time = AverageMeter()
data_time = AverageMeter()
end = time.time()
for idx, (inputs, _, index) in enumerate(train_loader):
data_time.update(time.time() - end)
# print(inputs[0].size())
bsz = inputs[0].size(0)
# fixed args.batch_size
for i in range(len(inputs)):
inputs[i] = inputs[i].float()
inputs[i] = inputs[i].cuda()
if bsz < opt.pt_batch_size:
print("batch less than 16, continue")
continue
index = index.cuda(non_blocking=True)
# ===================forward=====================
anchor, positive, negative = inputs
feat_n, _ = model(negative)
feat_q, _ = model(anchor)
feat_k, _ = model_ema(positive)
if feat_k.size(0) > feat_n.size(0):
print("wrong bsz")
out = contrast(feat_q, feat_k, feat_n, index)
prob = out[:, 0].mean()
loss = criterion(out)
# ===================backward=====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
# ===================meters=====================
loss_meter.update(loss.item(), bsz)
prob_meter.update(prob.item(), bsz)
moment_update(model, model_ema, opt.pt_alpha)
torch.cuda.synchronize()
batch_time.update(time.time() - end)
message = ('MoCo Train: [{0}][{1}/{2}]\t'
'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
'loss {loss.val:.3f} ({loss.avg:.3f})\t'
'prob {prob.val:.3f} ({prob.avg:.3f})'.format(
epoch,
idx + 1,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=loss_meter,
prob=prob_meter))
# print info
if (idx + 1) % opt.pt_print_freq == 0:
print(message)
recorder.record_message('a', message)
# print(out.shape)
sys.stdout.flush()
end = time.time()
return loss_meter.avg, prob_meter.avg
def test(epoch, cfg, data_loader, model, obj_vtx, obj_info, criterions):
model.eval()
Eval = Evaluation(cfg.dataset, obj_info, obj_vtx)
if 'trans' in cfg.pytorch.task.lower():
Eval_trans = Evaluation(cfg.dataset, obj_info, obj_vtx)
if not cfg.test.ignore_cache_file:
est_cache_file = cfg.test.cache_file
# gt_cache_file = cfg.test.cache_file.replace('pose_est', 'pose_gt')
gt_cache_file = cfg.test.cache_file.replace('_est', '_gt')
if os.path.exists(est_cache_file) and os.path.exists(gt_cache_file):
Eval.pose_est_all = np.load(est_cache_file,
allow_pickle=True).tolist()
Eval.pose_gt_all = np.load(gt_cache_file,
allow_pickle=True).tolist()
fig_save_path = os.path.join(cfg.pytorch.save_path, str(epoch))
mkdir_p(fig_save_path)
if 'all' in cfg.test.test_mode.lower():
Eval.evaluate_pose()
Eval.evaluate_pose_add(fig_save_path)
Eval.evaluate_pose_arp_2d(fig_save_path)
elif 'pose' in cfg.test.test_mode.lower():
Eval.evaluate_pose()
elif 'add' in cfg.test.test_mode.lower():
Eval.evaluate_pose_add(fig_save_path)
elif 'arp' in cfg.test.test_mode.lower():
Eval.evaluate_pose_arp_2d(fig_save_path)
else:
raise Exception("Wrong test mode: {}".format(
cfg.test.test_mode))
return None, None
else:
logger.info("test cache file {} and {} not exist!".format(
est_cache_file, gt_cache_file))
userAns = input("Generating cache file from model [Y(y)/N(n)]:")
if userAns.lower() == 'n':
sys.exit(0)
else:
logger.info("Generating test cache file!")
preds = {}
Loss = AverageMeter()
Loss_rot = AverageMeter()
Loss_trans = AverageMeter()
num_iters = len(data_loader)
bar = Bar('{}'.format(cfg.pytorch.exp_id[-60:]), max=num_iters)
time_monitor = False
vis_dir = os.path.join(cfg.pytorch.save_path, 'test_vis_{}'.format(epoch))
if not os.path.exists(vis_dir):
os.makedirs(vis_dir)
for i, (obj, obj_id, inp, pose, c_box, s_box, box,
trans_local) in enumerate(data_loader):
if cfg.pytorch.gpu > -1:
inp_var = inp.cuda(cfg.pytorch.gpu, async=True).float()
else:
inp_var = inp.float()
bs = len(inp)
# forward propagation
T_begin = time.time()
pred_rot, pred_trans = model(inp_var)
T_end = time.time() - T_begin
if time_monitor:
logger.info(
"time for a batch forward of resnet model is {}".format(T_end))
if i % cfg.test.disp_interval == 0:
# input image
inp_rgb = (inp[0].cpu().numpy().copy() *
255)[[2, 1, 0], :, :].astype(np.uint8)
cfg.writer.add_image('input_image', inp_rgb, i)
cv2.imwrite(os.path.join(vis_dir, '{}_inp.png'.format(i)),
inp_rgb.transpose(1, 2, 0)[:, :, ::-1])
if 'rot' in cfg.pytorch.task.lower():
# coordinates map
pred_coor = pred_rot[0, 0:3].data.cpu().numpy().copy()
pred_coor[0] = im_norm_255(pred_coor[0])
pred_coor[1] = im_norm_255(pred_coor[1])
pred_coor[2] = im_norm_255(pred_coor[2])
pred_coor = np.asarray(pred_coor, dtype=np.uint8)
cfg.writer.add_image('test_coor_x_pred',
np.expand_dims(pred_coor[0], axis=0), i)
cfg.writer.add_image('test_coor_y_pred',
np.expand_dims(pred_coor[1], axis=0), i)
cfg.writer.add_image('test_coor_z_pred',
np.expand_dims(pred_coor[2], axis=0), i)
# gt_coor = target[0, 0:3].data.cpu().numpy().copy()
# gt_coor[0] = im_norm_255(gt_coor[0])
# gt_coor[1] = im_norm_255(gt_coor[1])
# gt_coor[2] = im_norm_255(gt_coor[2])
# gt_coor = np.asarray(gt_coor, dtype=np.uint8)
# cfg.writer.add_image('test_coor_x_gt', np.expand_dims(gt_coor[0], axis=0), i)
# cfg.writer.add_image('test_coor_y_gt', np.expand_dims(gt_coor[1], axis=0), i)
# cfg.writer.add_image('test_coor_z_gt', np.expand_dims(gt_coor[2], axis=0), i)
# confidence map
pred_conf = pred_rot[0, 3].data.cpu().numpy().copy()
pred_conf = (im_norm_255(pred_conf)).astype(np.uint8)
cfg.writer.add_image('test_conf_pred',
np.expand_dims(pred_conf, axis=0), i)
# gt_conf = target[0, 3].data.cpu().numpy().copy()
# cfg.writer.add_image('test_conf_gt', np.expand_dims(gt_conf, axis=0), i)
if 'trans' in cfg.pytorch.task.lower():
pred_trans_ = pred_trans[0].data.cpu().numpy().copy()
gt_trans_ = trans_local[0].data.cpu().numpy().copy()
cfg.writer.add_scalar('test_trans_x_gt', gt_trans_[0],
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_y_gt', gt_trans_[1],
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_z_gt', gt_trans_[2],
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_x_pred', pred_trans_[0],
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_y_pred', pred_trans_[1],
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_z_pred', pred_trans_[2],
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_x_err',
np.abs(pred_trans_[0] - gt_trans_[0]),
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_y_err',
np.abs(pred_trans_[1] - gt_trans_[1]),
i + (epoch - 1) * num_iters)
cfg.writer.add_scalar('test_trans_z_err',
np.abs(pred_trans_[2] - gt_trans_[2]),
i + (epoch - 1) * num_iters)
if 'rot' in cfg.pytorch.task.lower():
pred_coor = pred_rot[:, 0:3].data.cpu().numpy().copy()
pred_conf = pred_rot[:, 3].data.cpu().numpy().copy()
else:
pred_coor = np.zeros(bs)
pred_conf = np.zeros(bs)
if 'trans' in cfg.pytorch.task.lower():
pred_trans = pred_trans.data.cpu().numpy().copy()
else:
pred_trans = np.zeros(bs)
col = list(
zip(obj, obj_id.numpy(), pred_coor, pred_conf, pred_trans,
pose.numpy(), c_box.numpy(), s_box.numpy(), box.numpy()))
for idx in range(len(col)):
obj_, obj_id_, pred_coor_, pred_conf_, pred_trans_, pose_gt, c_box_, s_box_, box_ = col[
idx]
T_begin = time.time()
if 'rot' in cfg.pytorch.task.lower():
# building 2D-3D correspondences
pred_coor_ = pred_coor_.transpose(1, 2, 0)
pred_coor_[:, :, 0] = pred_coor_[:, :, 0] * abs(
obj_info[obj_id_]['min_x'])
pred_coor_[:, :, 1] = pred_coor_[:, :, 1] * abs(
obj_info[obj_id_]['min_y'])
pred_coor_[:, :, 2] = pred_coor_[:, :, 2] * abs(
obj_info[obj_id_]['min_z'])
pred_coor_ = pred_coor_.tolist()
eroMask = False
if eroMask:
kernel = np.ones((3, 3), np.uint8)
pred_conf_ = cv2.erode(pred_conf_, kernel)
pred_conf_ = (pred_conf_ - pred_conf_.min()) / (
pred_conf_.max() - pred_conf_.min())
pred_conf_ = pred_conf_.tolist()
select_pts_2d = []
select_pts_3d = []
c_w = int(c_box_[0])
c_h = int(c_box_[1])
s = int(s_box_)
w_begin = c_w - s / 2.
h_begin = c_h - s / 2.
w_unit = s * 1.0 / cfg.dataiter.out_res
h_unit = s * 1.0 / cfg.dataiter.out_res
min_x = 0.001 * abs(obj_info[obj_id_]['min_x'])
min_y = 0.001 * abs(obj_info[obj_id_]['min_y'])
min_z = 0.001 * abs(obj_info[obj_id_]['min_z'])
for x in range(cfg.dataiter.out_res):
for y in range(cfg.dataiter.out_res):
if pred_conf_[x][y] < cfg.test.mask_threshold:
continue
if abs(pred_coor_[x][y][0]) < min_x and abs(pred_coor_[x][y][1]) < min_y and \
abs(pred_coor_[x][y][2]) < min_z:
continue
select_pts_2d.append(
[w_begin + y * w_unit, h_begin + x * h_unit])
select_pts_3d.append(pred_coor_[x][y])
model_points = np.asarray(select_pts_3d, dtype=np.float32)
image_points = np.asarray(select_pts_2d, dtype=np.float32)
if 'trans' in cfg.pytorch.task.lower():
# compute T from translation head
ratio_delta_c = pred_trans_[:2]
ratio_depth = pred_trans_[2]
pred_depth = ratio_depth * (cfg.dataiter.out_res / s_box_)
pred_c = ratio_delta_c * box_[2:] + c_box_
pred_x = (pred_c[0] - cfg.dataset.camera_matrix[0, 2]
) * pred_depth / cfg.dataset.camera_matrix[0, 0]
pred_y = (pred_c[1] - cfg.dataset.camera_matrix[1, 2]
) * pred_depth / cfg.dataset.camera_matrix[1, 1]
T_vector_trans = np.asarray([pred_x, pred_y, pred_depth])
pose_est_trans = np.concatenate(
(np.eye(3), np.asarray((T_vector_trans).reshape(3, 1))),
axis=1)
try:
if 'rot' in cfg.pytorch.task.lower():
dist_coeffs = np.zeros(
(4, 1)) # Assuming no lens distortion
if cfg.test.pnp == 'iterPnP': # iterative PnP algorithm
success, R_vector, T_vector = cv2.solvePnP(
model_points,
image_points,
cfg.dataset.camera_matrix,
dist_coeffs,
flags=cv2.SOLVEPNP_ITERATIVE)
elif cfg.test.pnp == 'ransac': # ransac algorithm
_, R_vector, T_vector, inliers = cv2.solvePnPRansac(
model_points,
image_points,
cfg.dataset.camera_matrix,
dist_coeffs,
flags=cv2.SOLVEPNP_EPNP)
else:
raise NotImplementedError(
"Not support PnP algorithm: {}".format(
cfg.test.pnp))
R_matrix = cv2.Rodrigues(R_vector, jacobian=0)[0]
pose_est = np.concatenate(
(R_matrix, np.asarray(T_vector).reshape(3, 1)), axis=1)
if 'trans' in cfg.pytorch.task.lower():
pose_est_trans = np.concatenate(
(R_matrix,
np.asarray((T_vector_trans).reshape(3, 1))),
axis=1)
Eval.pose_est_all[obj_].append(pose_est)
Eval.pose_gt_all[obj_].append(pose_gt)
Eval.num[obj_] += 1
Eval.numAll += 1
if 'trans' in cfg.pytorch.task.lower():
Eval_trans.pose_est_all[obj_].append(pose_est_trans)
Eval_trans.pose_gt_all[obj_].append(pose_gt)
Eval_trans.num[obj_] += 1
Eval_trans.numAll += 1
except:
Eval.num[obj_] += 1
Eval.numAll += 1
if 'trans' in cfg.pytorch.task.lower():
Eval_trans.num[obj_] += 1
Eval_trans.numAll += 1
logger.info('error in solve PnP or Ransac')
T_end = time.time() - T_begin
if time_monitor:
logger.info(
"time spend on PnP+RANSAC for one image is {}".format(
T_end))
Bar.suffix = 'test Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.4f} | Loss_rot {loss_rot.avg:.4f} | Loss_trans {loss_trans.avg:.4f}'.format(
epoch,
i,
num_iters,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=Loss,
loss_rot=Loss_rot,
loss_trans=Loss_trans)
bar.next()
epoch_save_path = os.path.join(cfg.pytorch.save_path, str(epoch))
if not os.path.exists(epoch_save_path):
os.makedirs(epoch_save_path)
if 'rot' in cfg.pytorch.task.lower():
logger.info("{} Evaluate of Rotation Branch of Epoch {} {}".format(
'-' * 40, epoch, '-' * 40))
preds['poseGT'] = Eval.pose_gt_all
preds['poseEst'] = Eval.pose_est_all
if cfg.pytorch.test:
np.save(os.path.join(epoch_save_path, 'pose_est_all_test.npy'),
Eval.pose_est_all)
np.save(os.path.join(epoch_save_path, 'pose_gt_all_test.npy'),
Eval.pose_gt_all)
else:
np.save(
os.path.join(epoch_save_path,
'pose_est_all_epoch{}.npy'.format(epoch)),
Eval.pose_est_all)
np.save(
os.path.join(epoch_save_path,
'pose_gt_all_epoch{}.npy'.format(epoch)),
Eval.pose_gt_all)
# evaluation
if 'all' in cfg.test.test_mode.lower():
Eval.evaluate_pose()
Eval.evaluate_pose_add(epoch_save_path)
Eval.evaluate_pose_arp_2d(epoch_save_path)
else:
if 'pose' in cfg.test.test_mode.lower():
Eval.evaluate_pose()
if 'add' in cfg.test.test_mode.lower():
Eval.evaluate_pose_add(epoch_save_path)
if 'arp' in cfg.test.test_mode.lower():
Eval.evaluate_pose_arp_2d(epoch_save_path)
if 'trans' in cfg.pytorch.task.lower():
logger.info("{} Evaluate of Translation Branch of Epoch {} {}".format(
'-' * 40, epoch, '-' * 40))
preds['poseGT'] = Eval_trans.pose_gt_all
preds['poseEst'] = Eval_trans.pose_est_all
if cfg.pytorch.test:
np.save(
os.path.join(epoch_save_path, 'pose_est_all_test_trans.npy'),
Eval_trans.pose_est_all)
np.save(
os.path.join(epoch_save_path, 'pose_gt_all_test_trans.npy'),
Eval_trans.pose_gt_all)
else:
np.save(
os.path.join(epoch_save_path,
'pose_est_all_trans_epoch{}.npy'.format(epoch)),
Eval_trans.pose_est_all)
np.save(
os.path.join(epoch_save_path,
'pose_gt_all_trans_epoch{}.npy'.format(epoch)),
Eval_trans.pose_gt_all)
# evaluation
if 'all' in cfg.test.test_mode.lower():
Eval_trans.evaluate_pose()
Eval_trans.evaluate_pose_add(epoch_save_path)
Eval_trans.evaluate_pose_arp_2d(epoch_save_path)
else:
if 'pose' in cfg.test.test_mode.lower():
Eval_trans.evaluate_pose()
if 'add' in cfg.test.test_mode.lower():
Eval_trans.evaluate_pose_add(epoch_save_path)
if 'arp' in cfg.test.test_mode.lower():
Eval_trans.evaluate_pose_arp_2d(epoch_save_path)
bar.finish()
return {
'Loss': Loss.avg,
'Loss_rot': Loss_rot.avg,
'Loss_trans': Loss_trans.avg
}, preds
def train_dsm_triplet(epoch, train_loader, model, optimizer, opt, pos_aug,
neg_aug, recorder):
"""
one epoch training for instance discrimination
"""
print("==> (DSM triplet) training...")
model.train()
def set_bn_train(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
m.train()
batch_time = AverageMeter()
data_time = AverageMeter()
loss_meter = AverageMeter()
triplet_loss = nn.TripletMarginLoss(margin=0.5, p=2)
end = time.time()
for idx, (inputs, _, index) in enumerate(train_loader):
data_time.update(time.time() - end)
bsz = inputs[0].size(0)
# fixed args.batch_size
if bsz < opt.pt_batch_size:
print("batch less than 16, continue")
continue
for i in range(len(inputs)):
inputs[i] = inputs[i].float()
inputs[i] = inputs[i].cuda()
# ===================forward=====================
anchor_old, positive, negative = inputs
# here a series of data augmentation
# ====================================================postive operation=======================
anchor = pos_aug(anchor_old)
feat_k = model(positive)
feat_n = model(negative)
feat_q = model(anchor)
if feat_k.size(0) > feat_q.size(0):
print("wrong bsz")
intra_loss = triplet_loss(feat_q, feat_k, feat_n)
inter_loss = triplet_loss(feat_q, feat_k, flip(feat_n, 0))
# for j in range(bsz-2):
# inter_loss += triplet_loss(feat_q, feat_k, shift(feat_n, 0))
alpha_1 = 1
alpha_2 = 1
loss = alpha_1 * intra_loss + alpha_2 * inter_loss
# print(loss)
# ===================backward=====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
# ===================meters=====================
loss_meter.update(loss.item(), bsz)
torch.cuda.synchronize()
batch_time.update(time.time() - end)
end = time.time()
message = ('DSM triplet Train: [{0}][{1}/{2}]\t'
'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
'loss {loss.val:.3f} ({loss.avg:.3f})'.format(
epoch,
idx + 1,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=loss_meter))
# print info
if (idx + 1) % opt.pt_print_freq == 0:
print(message)
recorder.record_message('a', message)
# print(out.shape)
sys.stdout.flush()
return loss_meter.avg
def train_model(epoch, model, optimizer, lr_scheduler, loader, test_loader):
global GLOBAL_STEP
test_loader_it = iter(test_loader)
loss_meter = AverageMeter()
val_loss_meter = AverageMeter()
acc_meter = AverageMeter()
val_acc_meter = AverageMeter()
model.train()
model.to(device)
print('=' * 20 + "Model Training" + '=' * 20)
loss_func = nn.CrossEntropyLoss()
for i, batch in tqdm(enumerate(loader)):
start = time.time()
model.train()
optimizer.zero_grad()
model.zero_grad()
sentence1, sentence2, label = batch
label = label.to(device)
pred = model((sentence1, sentence2))
loss = loss_func(pred, label)
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_bound)
optimizer.step()
acc = torch.mean((torch.max(pred, 1)[1] == label).type(torch.float))
loss_meter.update(loss.item())
acc_meter.update(acc.item())
end = time.time()
used_time = end - start
if (GLOBAL_STEP) % args.log_every == 0:
try:
batch = next(test_loader_it)
except:
test_loader_it = iter(test_loader)
batch = next(test_loader_it)
eval_loss, eval_acc = batch_eval(batch, model)
val_loss_meter.update(eval_loss.item())
val_acc_meter.update(eval_acc.item())
lr = optimizer.param_groups[0]['lr']
display = 'epoch=' + str(epoch) + \
'\tglobal_step=%d' % (GLOBAL_STEP) + \
'\tloss=%.4f' % (loss_meter.val) + \
'\tloss_avg=%.4f' % (loss_meter.avg) + \
'\tval_loss=%.4f' % (val_loss_meter.avg) + \
'\tacc=%.4f' % (acc_meter.avg) + \
'\tval_acc=%.4f' % (val_acc_meter.avg) + \
'\tlr=%.6f' % (lr) + \
'\t|g|=%.4f' % (grad_norm) + \
'\ttime=%.2fit/s' % (1. / used_time)
tqdm.write(display)
save_mode(epoch=epoch,
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler)
GLOBAL_STEP += 1
return
def train_dsm(epoch, train_loader, model, model_ema, contrast, criterion,
optimizer, opt, pos_aug, neg_aug, recorder):
"""
one epoch training for instance discrimination
"""
print("==> (DSM) training...")
model.train()
model_ema.eval()
def set_bn_train(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
m.train()
model_ema.apply(set_bn_train)
batch_time = AverageMeter()
data_time = AverageMeter()
loss_meter = AverageMeter()
prob_meter = AverageMeter()
end = time.time()
for idx, (inputs, _, index) in enumerate(train_loader):
data_time.update(time.time() - end)
# print(inputs[0].size())
bsz = inputs[0].size(0)
# fixed args.batch_size
if bsz < opt.pt_batch_size:
print("batch less than 16, continue")
continue
for i in range(len(inputs)):
inputs[i] = inputs[i].float()
inputs[i] = inputs[i].cuda()
index = index.cuda(non_blocking=True)
# ===================forward=====================
anchor_old, positive, negative = inputs
# print(anchor_old.size())
# here a series of data augmentation
# ====================================================postive operation=======================
anchor = pos_aug(anchor_old)
# positive = flip(anchor, 2)
# shuffle_ids, reverse_ids = get_shuffle_ids(bsz)
feat_q, map_q = model(anchor)
feat_k, map_k = model_ema(positive)
# tcr_loss = tcr(map_q, map_k)
# with torch.no_grad():
# positive = positive[shuffle_ids]
# feat_k = model_ema(positive)
# feat_k = feat_k[reverse_ids]
feat_n, _ = model(negative)
if feat_n.size(0) > feat_q.size(0):
print("wrong bsz")
out = contrast(feat_q, feat_k, feat_n, index)
contrast_loss = criterion(out)
loss = contrast_loss # + tcr_loss # + sample_loss # + contrast_loss2 # + cls_loss + mixup_loss
# print(contrast_loss, tcr_loss)
# ===================backward=====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
# ===================meters=====================
prob = out[:, 0].mean()
prob_meter.update(prob.item(), bsz)
loss_meter.update(loss.item(), bsz)
moment_update(model, model_ema, opt.pt_alpha)
torch.cuda.synchronize()
batch_time.update(time.time() - end)
end = time.time()
message = ('DSM Train: [{0}][{1}/{2}]\t'
'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
'loss {loss.val:.3f} ({loss.avg:.3f})\t'
'prob {prob.val:.3f} ({prob.avg:.3f})'.format(
epoch,
idx + 1,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=loss_meter,
prob=prob_meter))
# print info
if (idx + 1) % opt.pt_print_freq == 0:
print(message)
recorder.record_message('a', message)
# print(out.shape)
sys.stdout.flush()
return loss_meter.avg, prob_meter.avg
