def download_from_urls(urls, download_dir: str, verbosity):
"""
При отсутствии verbosity выводит кол-во и что скачивает как Прогресс бар
"""
if verbosity:
for url in urls:
download_url_content(url=url, download_dir=download_dir)
return None
bar = Bar('Processing', max=len(urls))
for url in urls:
bar.suffix = '%(index)d/%(max)d ' + url
bar.next()
download_url_content(url=url, download_dir=download_dir)
bar.finish()
def train(data_loader, model_pos, criterion, optimizer, device, lr_init, lr_now, step, decay, gamma, max_norm=True):
batch_time = AverageMeter()
data_time = AverageMeter()
epoch_loss_3d_pos = AverageMeter()
# Switch to train mode
torch.set_grad_enabled(True)
model_pos.train()
end = time.time()
bar = Bar('Train', max=len(data_loader))
for i, (targets_3d, inputs_2d, _) in enumerate(data_loader):
# Measure data loading time
data_time.update(time.time() - end)
num_poses = targets_3d.size(0)
step += 1
if step % decay == 0 or step == 1:
lr_now = lr_decay(optimizer, step, lr_init, decay, gamma)
targets_3d, inputs_2d = targets_3d.to(device), inputs_2d.to(device)
"""
inputs_2d : 模型的输入参数 =========================================================================================《《《
targets_3d : 目标结果
"""
# print(inputs_2d[0])
# print(inputs_2d.shape)
outputs_3d = model_pos(inputs_2d) # inputs_2d到model中得到outputs_3d============================================
optimizer.zero_grad()
loss_3d_pos = criterion(outputs_3d, targets_3d)
loss_3d_pos.backward()
if max_norm:
nn.utils.clip_grad_norm_(model_pos.parameters(), max_norm=1)
optimizer.step()
epoch_loss_3d_pos.update(loss_3d_pos.item(), num_poses)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {ttl:} | ETA: {eta:} ' \
'| Loss: {loss: .4f}' \
.format(batch=i + 1, size=len(data_loader), data=data_time.avg, bt=batch_time.avg,
ttl=bar.elapsed_td, eta=bar.eta_td, loss=epoch_loss_3d_pos.avg)
bar.next()
bar.finish()
return epoch_loss_3d_pos.avg, lr_now, step
def train(train_loader, model, criterion, optimizer, epoch, device):
# 模型进入训练模式
model.train()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
end = time.time()
bar = Bar('Processing train', max=len(train_loader))
for i, (data, target) in enumerate(train_loader):
# 计算加载数据的时间
data_time.update(time.time() - end)
# 备注: 分布式训练中要设置non_blocking=True
data = data.to(device)
target = target.to(device)
output = model(data)
loss = criterion(output, target)
prec1, prec5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), data.size(0))
top1.update(prec1[0], data.size(0))
top5.update(prec5[0], data.size(0))
# 计算一个批量计算的时间
batch_time.update(time.time() - end)
end = time.time()
optimizer.zero_grad()
loss.backward()
optimizer.step()
bar.suffix = '(Epoch:{epoch} - {batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Loss: {loss:.2f} | Acc@1: {acc1:.2f} | Acc@5: {acc5:.2f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.avg,
bt=batch_time.avg,
loss=losses.avg,
acc1=top1.avg,
acc5=top5.avg,
epoch=epoch)
bar.next()
bar.finish()
return losses.avg, top1.avg, top5.avg
def test(test_loader, model, criterion, use_cuda):
# switch to evaluate mode
model.eval()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
end = time.time()
bar = Bar('Processing', max=len(test_loader))
for batch_idx, (inputs, targets) in enumerate(test_loader):
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
with torch.no_grad():
# compute output
outputs = model(inputs)
loss = criterion(outputs, targets) / args.num_partitionings
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))
losses.update(loss.data.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(
batch=batch_idx + 1,
size=len(test_loader),
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
top1=top1.avg,
top5=top5.avg,
)
bar.next()
bar.finish()
return losses.avg, top1.avg
def train(model, train_loader, epoch, criterion, optimizer, writer):
# set training mode
model.train()
train_loss = 0.0
iter_num = 0
# Iterate over data.
bar = Bar('Processing | {}'.format('train'), max=len(train_loader))
bar.check_tty = False
for i_iter, batch in enumerate(train_loader):
sys.stdout.flush()
start_time = time.time()
iter_num += 1
# adjust learning rate
iters_per_epoch = len(train_loader)
lr = adjust_learning_rate(optimizer, epoch, i_iter, iters_per_epoch, method=args.lr_mode)
image, label, hlabel, flabel, _ = batch
images, labels, hlabel, flabel = image.cuda(), label.long().cuda(), hlabel.cuda(), flabel.cuda()
torch.set_grad_enabled(True)
# zero the parameter gradients
optimizer.zero_grad()
# compute output loss
preds = model(images)
loss = criterion(preds, [labels, hlabel, flabel]) # batch mean
train_loss += loss.item()
# compute gradient and do SGD step
loss.backward()
optimizer.step()
if i_iter % 10 == 0:
writer.add_scalar('learning_rate', lr, iter_num + epoch * len(train_loader))
writer.add_scalar('train_loss', train_loss / iter_num, iter_num + epoch * len(train_loader))
batch_time = time.time() - start_time
# plot progress
bar.suffix = '{} / {} | Time: {batch_time:.4f} | Loss: {loss:.4f}'.format(iter_num, len(train_loader),
batch_time=batch_time,
loss=train_loss / iter_num)
bar.next()
epoch_loss = train_loss / iter_num
writer.add_scalar('train_epoch_loss', epoch_loss, epoch)
bar.finish()
return epoch_loss
def train(train_loader, model, criterion, optimizer, args):
batch_time = AverageMeter()
data_time = AverageMeter()
am_shape_l2 = AverageMeter()
am_bone_len = AverageMeter()
last = time.time()
# switch to trian
model.train()
bar = Bar('\033[31m Train \033[0m', max=len(train_loader))
for i, metas in enumerate(train_loader):
data_time.update(time.time() - last)
results, targets, total_loss, losses = one_forward_pass(
metas, model, criterion, args, train=True
)
global steps
steps += 1
writer.add_scalar('loss', total_loss.item(), steps)
am_shape_l2.update(losses['shape_reg'].item(), targets['batch_size'])
am_bone_len.update(losses['bone_len'].item(), targets['batch_size'])
''' backward and step '''
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
''' progress '''
batch_time.update(time.time() - last)
last = time.time()
bar.suffix = (
'({batch}/{size}) '
'd: {data:.2f}s | '
'b: {bt:.2f}s | '
't: {total:}s | '
'eta:{eta:}s | '
'lN: {lossLen:.5f} | '
'lL2: {lossL2:.5f} | '
).format(
batch=i + 1,
size=len(train_loader),
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
lossLen=am_bone_len.avg,
lossL2=am_shape_l2.avg,
)
bar.next()
bar.finish()
return (am_bone_len.avg, am_shape_l2.avg)
def train_one_epoch(epoch_number):
bar = Bar('Processing', max=len(train_loader))
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
accuracy = AverageMeter()
model.train()
end = time.time()
for i, (x, y) in enumerate(train_loader):
data_time.update(time.time() - end)
x, y = x.cuda(), y.cuda()
output = model(x)
if not ALL_CLASSES:
output = output.view(-1)
y = y[:, CLASS]
loss = loss_criterion(output, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
y_hat = output > 0.5
correct = y_hat == y.type(torch.cuda.BoolTensor)
accuracy.update(correct.sum().item() / correct.numel())
losses.update(loss.item())
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = "({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.9f} | Acc: {acc:.2f}%%".format(
batch=i + 1,
size=len(train_loader),
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=accuracy.avg * 100)
bar.next()
bar.finish()
return losses.avg
def _img_crop(self, img, img_name, i, tp):
_, height, width = img.shape
len_y, len_x = self.crop_size
num_imgs = (height // (len_x // 2) + 1) * (width // (len_y // 2) + 1)
bar = Bar(f'{tp} {i + 1} spliting:', max=num_imgs)
x = 0
row_count = 0
while x < height:
y = 0
col_count = 0
while y < width:
if y > width - len_y and x > height - len_x:
split_image = img[:, height - len_x:, width - len_y:]
elif y > width - len_y:
split_image = img[:, x:x + len_x, width - len_y:]
elif x > height - len_x:
split_image = img[:, height - len_x:, y:y + len_y]
else:
split_image = img[:, x:x + len_x, y:y + len_y]
split_image_name = '_'.join(
[img_name, str(row_count),
str(col_count)])
if tp == 'image':
make_sure_path_exists(os.path.join(self.save_path, 'img'))
np.save(
os.path.join(self.save_path, 'img', split_image_name),
split_image)
else:
make_sure_path_exists(os.path.join(self.save_path,
'label'))
# split_image_name = '_'.join([img_name.replace('_label', ''), str(row_count), str(col_count)])
np.save(
os.path.join(self.save_path, 'label',
split_image_name), split_image)
y = min(width, y + len_y // 2)
col_count += 1
bar.suffix = f'{row_count * (width // (len_y // 2) + 1) + col_count}/{num_imgs}'
bar.next()
if x == height:
break
x = min(height, x + len_x // 2)
row_count += 1
bar.finish()
def validate(valloader, model, criterion, epoch, use_cuda, mode):
batch_time = AverageMeter("batch_time")
data_time = AverageMeter("data_time")
losses = AverageMeter("loss")
top1 = AverageMeter("top1")
top5 = AverageMeter("top5")
# switch to evaluate mode
model.eval()
end = time.time()
bar = Bar(f'{mode}', max=len(valloader))
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(valloader):
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda(
non_blocking=True)
# compute output
outputs = model(inputs)
loss = criterion(outputs, targets)
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs, targets, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(
batch=batch_idx + 1,
size=len(valloader),
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
top1=top1.avg,
top5=top5.avg,
)
bar.next()
bar.finish()
return (losses.avg, top1.avg)
def statistic(data_path):
data_list = os.listdir(os.path.join(data_path, 'mask'))
num = len(data_list)
bar = Bar('counting:', max=num)
res = np.zeros(16)
for idx, data_file in enumerate(data_list):
mask = np.load(os.path.join(data_path, 'mask', data_file))
for i in range(16):
count = (mask == i).sum()
res[i] += count
bar.suffix = '{} / {}'.format(idx, num)
bar.next()
bar.finish()
return res
def _std(path, img_list, mean, pixels_num):
files_num = len(img_list)
bar = Bar('Calculating std:', max=files_num)
value_std = [0, 0, 0, 0]
i = 0
for img_file in img_list:
img = np.load(os.path.join(path, img_file)) / 255.0
value_std += np.sum((img.transpose(
(1, 2, 0)) - mean).transpose(2, 0, 1)**2,
axis=(1, 2))
i += 1
bar.suffix = f'{i}/{files_num}'
bar.next()
bar.finish()
return np.sqrt(value_std / pixels_num)
def evaluate(data_loader, model_pos, device):
batch_time = AverageMeter()
data_time = AverageMeter()
epoch_loss_3d_pos = AverageMeter()
epoch_loss_3d_pos_procrustes = AverageMeter()
# Switch to evaluate mode
torch.set_grad_enabled(False)
model_pos.eval()
end = time.time()
bar = Bar('Eval ', max=len(data_loader))
for i, (targets_3d, inputs_2d, _) in enumerate(data_loader):
# Measure data loading time
data_time.update(time.time() - end)
num_poses = targets_3d.size(0)
inputs_2d = inputs_2d.to(device)
outputs_3d = model_pos(inputs_2d).cpu()
outputs_3d[:, :, :] -= outputs_3d[:, :
1, :] # Zero-centre the root (hip)
### remove scale
#inputs_2d = inputs_2d[:,:,:] - inputs_2d[:,:1,:]
#scale = compute_scale_pytorch(inputs_2d.cpu(), targets_3d)
#scale = torch.unsqueeze(torch.unsqueeze(scale, dim=1).expand(scale.size(0),17),dim=2)
#outputs_3d = outputs_3d*scale
###
epoch_loss_3d_pos.update(
mpjpe(outputs_3d, targets_3d).item() * 1000.0, num_poses)
epoch_loss_3d_pos_procrustes.update(
p_mpjpe(outputs_3d.numpy(), targets_3d.numpy()).item() * 1000.0,
num_poses)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {ttl:} | ETA: {eta:} ' \
'| MPJPE: {e1: .4f} | P-MPJPE: {e2: .4f}' \
.format(batch=i + 1, size=len(data_loader), data=data_time.val, bt=batch_time.avg,
ttl=bar.elapsed_td, eta=bar.eta_td, e1=epoch_loss_3d_pos.avg, e2=epoch_loss_3d_pos_procrustes.avg)
bar.next()
bar.finish()
return epoch_loss_3d_pos.avg, epoch_loss_3d_pos_procrustes.avg
def test(testloader, model, criterion, epoch, use_cuda):
global best_acc
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
bar = Bar('Processing', max=len(testloader))
for batch_idx, (inputs, targets) in enumerate(testloader):
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
inputs, targets = torch.autograd.Variable(
inputs, volatile=True), torch.autograd.Variable(targets)
# compute output
outputs = model(inputs)
loss = criterion(outputs, targets)
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size})| Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(
batch=batch_idx + 1,
size=len(testloader),
loss=losses.avg,
top1=top1.avg,
top5=top5.avg,
)
bar.next()
bar.finish()
return (losses.avg, top1.avg)
def train(train_loader, model, criterion, optimizer, scheduler, device,
config):
# switch to train mode
model.train()
losses = AverageMeter()
dice = AverageMeter()
bar = Bar('Processing', max=len(train_loader))
for step, (inputs, targets) in enumerate(train_loader):
images = inputs.cuda(device, non_blocking=True)
targets = targets.cuda(device, non_blocking=True)
outputs = model(images).squeeze(1)
loss = criterion(outputs, targets)
# measure roc auc and record loss
losses.update(loss.item(), images.size(0))
pred_probs = torch.sigmoid(outputs)
dices = general_dice(
targets.detach().cpu().numpy(),
torch.argmax(pred_probs, 1).detach().cpu().numpy())
dice.update(dices, images.size(0))
optimizer.zero_grad()
loss.backward()
# compute gradient and do SGD step
optimizer.step()
scheduler.step()
# plot progress
bar.suffix = '({batch}/{size})| LR: {lr:} | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Dice: {mean_dice: .4f}'.format(
batch=step + 1,
size=len(train_loader),
lr=get_lr(optimizer, -1),
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
mean_dice=dice.avg)
bar.next()
bar.finish()
return losses.avg
def resize_img(img_path, save_path, size=(12, 12)):
file_list = os.listdir(img_path)
num_imgs = len(file_list)
bar = Bar('resizing images:', max=num_imgs)
make_sure_path_exists(save_path)
for idx, file in enumerate(file_list):
img = Image.fromarray(np.load(os.path.join(img_path, file)))
img = np.array(img.resize(size))
np.save(os.path.join(save_path, file), img)
bar.suffix = f'{idx + 1} / {num_imgs}'
bar.next()
bar.finish()
def train_model(model, dataset, criterion, optimizer, epochs, dumping_dir):
model.train()
dataloader = DataLoader(dataset, shuffle=True, batch_size=4)
for i in range(epochs):
bar = Bar(f"Epoch {i + 1}", max=len(dataset) // 4)
for batch_idx, (src_tensor, tgt_tensor) in enumerate(iter(dataloader)):
optimizer.zero_grad()
output = model(src_tensor)
loss = criterion(output, tgt_tensor)
loss.backward()
optimizer.step()
bar.suffix = f"Loss {loss.item()}"
bar.next()
torch.save(model.state_dict(), os.path.join(dumping_dir, f"epoch_{i}.pt"))
return model
def distributing(path):
img_list = os.listdir(path)
num_imgs = len(img_list)
bar = Bar('distributing:', max=num_imgs)
dis = np.zeros(NUM_CLASSES)
for i, mask_file in enumerate(img_list):
mask = np.load(os.path.join(path, mask_file))
for category in range(NUM_CLASSES):
dis[category] += (mask == category).sum()
# print(ratios[category])
bar.suffix = f'{i + 1} / {num_imgs}'
bar.next()
bar.finish()
print(dis)
def train(train_loader,
model,
optimizer,
lr_now=None,
max_norm=True,
is_cuda=False,
dct_n=15,
dim_used=[]):
t_3d = utils.AccumLoss()
model.train()
st = time.time()
bar = Bar('>>>', fill='>', max=len(train_loader))
for i, (inputs, targets, all_seq) in enumerate(train_loader):
batch_size = inputs.shape[0]
if batch_size == 1:
break
bt = time.time()
if is_cuda:
inputs = Variable(inputs.cuda()).float()
# targets = Variable(targets.cuda(async=True)).float()
all_seq = Variable(all_seq.cuda(async=True)).float()
else:
inputs = Variable(inputs).float()
# targets = Variable(targets).float()
all_seq = Variable(all_seq).float()
outputs = model(inputs)
m_err = loss_funcs.mpjpe_error_3dpw(outputs, all_seq, dct_n, dim_used)
# calculate loss and backward
optimizer.zero_grad()
m_err.backward()
if max_norm:
nn.utils.clip_grad_norm(model.parameters(), max_norm=1)
optimizer.step()
n, seq_len, _ = all_seq.data.shape
t_3d.update(m_err.cpu().data.numpy()[0] * n * seq_len, n * seq_len)
bar.suffix = '{}/{}|batch time {:.4f}s|total time{:.2f}s'.format(
i, len(train_loader),
time.time() - bt,
time.time() - st)
bar.next()
bar.finish()
return lr_now, t_3d.avg
def evaluate(data_loader, model_pos, device):
batch_time = AverageMeter()
data_time = AverageMeter()
epoch_loss_3d_pos = AverageMeter()
epoch_loss_3d_pos_procrustes = AverageMeter()
# Switch to evaluate mode
torch.set_grad_enabled(False)
model_pos.eval()
end = time.time()
bar = Bar('Eval ', max=len(data_loader))
for i, (targets_3d, inputs_2d, _, intrinsics) in enumerate(data_loader):
# Measure data loading time
data_time.update(time.time() - end)
num_poses = targets_3d.size(0)
inputs_2d = inputs_2d.to(device)
outputs_3d = model_pos(inputs_2d).cpu()
output_2d = project_to_2d_linear(outputs_3d, intrinsics)
# true_output_2d = project_to_2d(targets_3d, intrinsics)
true_output_2d = project_to_2d_linear(targets_3d, intrinsics)
outputs_3d[:, :, :] -= outputs_3d[:, :
1, :] # Zero-centre the root (hip)
targets_3d[:, :, :] -= targets_3d[:, :
1, :] # Zero-centre the root (hip)
epoch_loss_3d_pos.update(
mpjpe(outputs_3d, targets_3d).item() * 1000.0, num_poses)
epoch_loss_3d_pos_procrustes.update(
p_mpjpe(outputs_3d.numpy(), targets_3d.numpy()).item() * 1000.0,
num_poses)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {ttl:} | ETA: {eta:} ' \
'| MPJPE: {e1: .4f} | P-MPJPE: {e2: .4f}' \
.format(batch=i + 1, size=len(data_loader), data=data_time.val, bt=batch_time.avg,
ttl=bar.elapsed_td, eta=bar.eta_td, e1=epoch_loss_3d_pos.avg, e2=epoch_loss_3d_pos_procrustes.avg)
bar.next()
bar.finish()
return epoch_loss_3d_pos.avg, epoch_loss_3d_pos_procrustes.avg
def validate(self):
self.generator.eval()
start = time.time()
summary_string = ''
bar = Bar('Validation', fill='#', max=len(self.valid_loader))
if self.evaluation_accumulators is not None:
for k, v in self.evaluation_accumulators.items():
self.evaluation_accumulators[k] = []
J_regressor = torch.from_numpy(
np.load(osp.join(VIBE_DATA_DIR, 'J_regressor_h36m.npy'))).float()
for i, target in enumerate(self.valid_loader):
move_dict_to_device(target, self.device)
# <=============
with torch.no_grad():
inp = target['features']
preds = self.generator(inp, J_regressor=J_regressor)
# convert to 14 keypoint format for evaluation
n_kp = preds[-1]['kp_3d'].shape[-2]
pred_j3d = preds[-1]['kp_3d'].view(-1, n_kp, 3).cpu().numpy()
target_j3d = target['kp_3d'].view(-1, n_kp, 3).cpu().numpy()
pred_verts = preds[-1]['verts'].view(-1, 6890, 3).cpu().numpy()
target_theta = target['theta'].view(-1, 85).cpu().numpy()
self.evaluation_accumulators['pred_verts'].append(pred_verts)
self.evaluation_accumulators['target_theta'].append(
target_theta)
self.evaluation_accumulators['pred_j3d'].append(pred_j3d)
self.evaluation_accumulators['target_j3d'].append(target_j3d)
# =============>
batch_time = time.time() - start
summary_string = f'({i + 1}/{len(self.valid_loader)}) | batch: {batch_time * 10.0:.4}ms | ' \
f'Total: {bar.elapsed_td} | ETA: {bar.eta_td:}'
bar.suffix = summary_string
bar.next()
bar.finish()
logger.info(summary_string)
def _val_one_epoch(self):
bar = Bar('Validating', max=len(self.val_data))
self.model.eval()
predictions = np.empty((0, 2))
targetcoors = np.empty((0, 2))
for step, (data, label) in enumerate(self.val_data):
with t.no_grad():
inputs = data
target = label.reshape(-1, 2)
# target = label.type(t.FloatTensor)
if len(self.params.gpus) > 0:
inputs = inputs.cuda(1)
# target = target.cuda()
score = self.model(inputs)
coors = spatial_soft_argmax2d(score[len(score) - 1], 1000,
False).cpu().numpy().reshape(-1, 2)
predictions = np.concatenate((predictions, coors), axis=0)
targetcoors = np.concatenate((targetcoors, target), axis=0)
# evaluation: calculate PCKh
currentpckh, currenterr = evalPCKh(predictions,
targetcoors,
threshold=50,
alpha=0.2)
# tensorboard visualization
if step % 100 == 0:
self.writer.add_image('valid/img', inputs[0], self.last_epoch)
self.writer.add_image('valid/output',
t.sum(score[1][0], 0),
self.last_epoch,
dataformats='HW')
bar.suffix = 'Valid: [%(index)d/%(max)d] | PCKh: {pckh:6f} | AveErr: {err:.2f} pixel |'.format(
pckh=currentpckh, err=currenterr)
bar.next()
bar.finish()
self.PCkh, self.AveErr = evalPCKh(predictions,
targetcoors,
threshold=50,
alpha=0.2)
self.model.train()
def train_autoencoder(model, dl, args, device, model_path):
model.train()
# start = time.time()
# threshold = 0.1
criterion = nn.BCELoss().to(device)
# optimizer = torch.optim.SGD(
# model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.00001)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.ae_lr,
weight_decay=0.00001)
def lr_lambda_ae(epoch):
if epoch < 50:
return 1
elif epoch >= 50 and epoch < 70:
return 0.16666667
elif epoch >= 70 and epoch < 90:
return 0.0335
else:
return 0.01
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda_ae)
for ep in range(args.nb_epoch_ae):
train_loss_list = []
bar = Bar(f"AE Training, epoch: {ep}", max=len(dl))
for i, (inputs, _) in enumerate(dl):
inputs = inputs.to(device)
_, outputs = model(inputs)
loss = criterion(outputs, inputs)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
train_loss_list.append(loss.item())
bar.suffix = (
f"({i}/{len(dl)}) | ETA: {bar.eta_td} | Loss: {loss.item():.4f} | "
f"Avg. Loss: {np.mean(np.array(train_loss_list)):.4f}")
bar.next()
bar.finish()
print("Saving AE model..")
state = {"model": model.encoder.state_dict(), "acc": 0, "epoch": -1}
torch.save(state, model_path)
def fit_pca():
pca = IncrementalPCA(n_components=PCA_COMPONENTS)
loader = get_fake_loader(FIT_PCA_BATCH_SIZE, optional_flip=False)
N_BATCHES = min(len(loader), ceil(N_DATA_POINTS_PCA / FIT_PCA_BATCH_SIZE))
bar = Bar('Fitting', max=N_BATCHES)
start = time.time()
for i, batch in enumerate(loader):
print(i)
if i >= N_BATCHES:
break
X = batch.view(FIT_PCA_BATCH_SIZE, -1)
pca.partial_fit(X)
bar.suffix = f"({(i+1)*FIT_PCA_BATCH_SIZE}/{N_BATCHES*FIT_PCA_BATCH_SIZE}) | Time: {time.time()-start:.1f}s | "
bar.next()
components = pca.mean_, pca.components_
pickle.dump(components, open(WEIGHTS_PATH, 'wb'))
bar.finish()
def validate(valloader, model, criterion, use_cuda, mode):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
Acc = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
bar = Bar(f'{mode}', max=len(valloader))
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(valloader):
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda(non_blocking=True)
# compute output
outputs = model(inputs)
# 注意和训练的时候计算方式有些不一样
loss = criterion(outputs[:, 1], targets)
# measure accuracy and record loss
acc = accuracy(outputs, targets)
losses.update(loss.item(), inputs.size(0))
Acc.update(acc.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=batch_idx + 1,
size=len(valloader),
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=Acc.avg,
)
bar.next()
bar.finish()
return (losses.avg, Acc.avg)
def train_posenet(model_pos, data_loader, optimizer, criterion, device):
batch_time = AverageMeter()
data_time = AverageMeter()
epoch_loss_3d_pos = AverageMeter()
# Switch to train mode
torch.set_grad_enabled(True)
set_grad([model_pos], True)
model_pos.train()
end = time.time()
bar = Bar('Train posenet', max=len(data_loader))
for i, (targets_3d, inputs_2d, _, _) in enumerate(data_loader):
# Measure data loading time
data_time.update(time.time() - end)
num_poses = targets_3d.size(0)
# here avoid bn with one sample in last batch, skip if num_poses=1
if num_poses == 1:
break
targets_3d, inputs_2d = targets_3d.to(device), inputs_2d.to(device)
targets_3d = targets_3d[:, :, :] - targets_3d[:, :
1, :] # the output is relative to the 0 joint
outputs_3d = model_pos(inputs_2d)
optimizer.zero_grad()
loss_3d_pos = criterion(outputs_3d, targets_3d)
loss_3d_pos.backward()
nn.utils.clip_grad_norm_(model_pos.parameters(), max_norm=1)
optimizer.step()
epoch_loss_3d_pos.update(loss_3d_pos.item(), num_poses)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {ttl:} | ETA: {eta:} ' \
'| Loss: {loss: .4f}' \
.format(batch=i + 1, size=len(data_loader), data=data_time.avg, bt=batch_time.avg,
ttl=bar.elapsed_td, eta=bar.eta_td, loss=epoch_loss_3d_pos.avg)
bar.next()
bar.finish()
return
def train_valid(paths_dict):
_data_list = os.listdir(os.path.join(paths_dict['source_path'], 'img'))
num_files = len(_data_list)
num_train = int(num_files * 0.9)
num_val = num_files - num_train
bar = Bar('Dividing trainset and validset:', max=num_files)
for i in range(num_files):
file = random.choice(_data_list)
name = file.split(".")[0]
_data_list.remove(file)
img_source = os.path.join(paths_dict['source_path'], 'img', file)
label_source = os.path.join(paths_dict['source_path'], 'label', file)
mask_source = os.path.join(paths_dict['source_path'], 'mask', file)
sta_source = os.path.join(paths_dict['source_path'], 'ratios', file)
if i < num_train:
img_target = make_sure_path_exists(
os.path.join(paths_dict['tr_save_path'], 'img'))
label_target = make_sure_path_exists(
os.path.join(paths_dict['tr_save_path'], 'label'))
mask_target = make_sure_path_exists(
os.path.join(paths_dict['tr_save_path'], 'mask'))
sta_target = make_sure_path_exists(
os.path.join(paths_dict['tr_save_path'], 'ratios'))
else:
img_target = make_sure_path_exists(
os.path.join(paths_dict['vd_save_path'], 'img'))
label_target = make_sure_path_exists(
os.path.join(paths_dict['vd_save_path'], 'label'))
mask_target = make_sure_path_exists(
os.path.join(paths_dict['vd_save_path'], 'mask'))
sta_target = make_sure_path_exists(
os.path.join(paths_dict['vd_save_path'], 'ratios'))
shutil.copy(img_source, img_target)
shutil.copy(label_source, label_target)
shutil.copy(mask_source, mask_target)
shutil.copy(sta_source, sta_target)
bar.suffix = f'{i + 1}/{num_files}'
bar.next()
bar.finish()
def val(train_loader,
model,
input_n=20,
output_n=10,
is_cuda=False,
dim_used=[]):
# t_l = utils.AccumLoss()
t_e = utils.AccumLoss()
t_3d = utils.AccumLoss()
model.eval()
st = time.time()
bar = Bar('>>>', fill='>', max=len(train_loader))
for i, (inputs, targets, all_seq) in enumerate(train_loader):
bt = time.time()
if is_cuda:
inputs = Variable(inputs.cuda()).float()
# targets = Variable(targets.cuda(async=True)).float()
all_seq = Variable(all_seq.cuda()).float()
outputs = model(inputs)
n = outputs.shape[0]
outputs = outputs.view(n, -1)
# targets = targets.view(n, -1)
# loss = loss_funcs.sen_loss(outputs, all_seq, dim_used)
n, _, _ = all_seq.data.shape
m_err = loss_funcs.mpjpe_error(outputs, all_seq, input_n, dim_used,
input_n + output_n)
e_err = loss_funcs.euler_error(outputs, all_seq, input_n, dim_used,
input_n + output_n)
# t_l.update(loss.cpu().data.numpy()[0] * n, n)
t_e.update(e_err.cpu().data.numpy() * n, n)
t_3d.update(m_err.cpu().data.numpy() * n, n)
bar.suffix = '{}/{}|batch time {:.4f}s|total time{:.2f}s'.format(
i + 1, len(train_loader),
time.time() - bt,
time.time() - st)
bar.next()
bar.finish()
return t_e.avg, t_3d.avg
def evaluate(data_loader, model_pos, device, architecture):
batch_time = AverageMeter()
data_time = AverageMeter()
epoch_loss_3d_pos = AverageMeter()
epoch_loss_3d_pos_procrustes = AverageMeter()
predictions = []
# Switch to evaluate mode
torch.set_grad_enabled(False)
model_pos.eval()
end = time.time()
bar = Bar('Eval ', max=len(data_loader))
for i, (targets_3d, inputs_2d, _) in enumerate(data_loader):
# Measure data loading time
data_time.update(time.time() - end)
num_poses = targets_3d.size(0)
inputs_2d = inputs_2d.to(device)
if architecture == 'linear':
outputs_3d = model_pos(inputs_2d.view(num_poses, -1)).view(num_poses, -1, 3).cpu()
outputs_3d = torch.cat([torch.zeros(num_poses, 1, outputs_3d.size(2)), outputs_3d], 1) # Pad hip joint
else:
outputs_3d = model_pos(inputs_2d).cpu()
outputs_3d[:, :, :] -= outputs_3d[:, :1, :] # Zero-centre the root (hip)
predictions.append(outputs_3d.numpy())
epoch_loss_3d_pos.update(mpjpe(outputs_3d, targets_3d).item() * 1000.0, num_poses)
epoch_loss_3d_pos_procrustes.update(p_mpjpe(outputs_3d.numpy(), targets_3d.numpy()).item() * 1000.0, num_poses)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {ttl:} | ETA: {eta:} ' \
'| MPJPE: {e1: .4f} | P-MPJPE: {e2: .4f}' \
.format(batch=i + 1, size=len(data_loader), data=data_time.val, bt=batch_time.avg,
ttl=bar.elapsed_td, eta=bar.eta_td, e1=epoch_loss_3d_pos.avg, e2=epoch_loss_3d_pos_procrustes.avg)
bar.next()
bar.finish()
return np.concatenate(predictions), epoch_loss_3d_pos.avg, epoch_loss_3d_pos_procrustes.avg
def test_avs(valloader, model, cluster_a, cluster_v, epoch, use_gpu):
model.eval()
data_time = AverageMeter()
match_rate = AverageMeter()
differ_rate = AverageMeter()
end = time.time()
bar = Bar('Processing', max=len(valloader))
for batch_idx, (audio, visual, roi) in enumerate(valloader):
audio = audio.view(args.val_batch * args.mix, *audio.shape[-2:])
visual = visual.view(args.val_batch * args.mix * args.frame, *visual.shape[-3:])
roi = roi.view(args.val_batch, args.mix * args.frame, args.rois, 4)
data_time.update(time.time() - end)
if use_gpu:
audio = audio.cuda()
visual = visual.cuda()
roi = roi.cuda()
data_time.update(time.time() - end)
discrim, mask, pred_a, pred_v, label_a, label_v, _, _ = model(audio, visual, roi, cluster_a, cluster_v)
discrim = [discrim[0].view(-1, 1), discrim[1].view(-1, 1)]
discrim = torch.cat(discrim, 0)
segments = discrim.shape[0]
true_match = torch.sum(discrim[: segments//2, 0]<1.0)
true_match = true_match.item() / float(segments/2)
true_differ = torch.sum(discrim[segments//2:, 0]>1.0)
true_differ = true_differ.item() / float(segments/2)
match_rate.update(true_match, 1)
differ_rate.update(true_differ, 1)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.3f}s |Match: {match:.3f} |Differ: {differ: .3f}'.format(
batch=batch_idx + 1,
size=len(valloader),
data=data_time.val,
match=match_rate.val,
differ=differ_rate.val
)
bar.next()
bar.finish()
return match_rate.avg, differ_rate.avg
def mean_std(path):
img_list = os.listdir(path)
pixels_num = 0
value_sum = [0, 0, 0]
files_num = len(img_list)
bar = Bar('Calculating mean:', max=files_num)
for idx, img_file in enumerate(img_list):
img = np.load(os.path.join(path, img_file)) / 255.0
pixels_num += img.shape[0] * img.shape[1]
value_sum += np.sum(img, axis=(0, 1))
bar.suffix = f'{idx + 1} / {files_num}'
bar.next()
bar.finish()
value_mean = value_sum / pixels_num
value_std = _std(path, img_list, value_mean, pixels_num)
return value_mean, value_std
def validate(loader, model, criterion, netType, debug, flip):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
end = time.time()
# predictions
predictions = torch.Tensor(loader.dataset.__len__(), 68, 2)
model.eval()
gt_win, pred_win = None, None
bar = Bar('Validating', max=len(loader))
all_dists = torch.zeros((68, loader.dataset.__len__()))
for i, (inputs, target, meta) in enumerate(loader):
data_time.update(time.time() - end)
input_var = torch.autograd.Variable(inputs.cuda())
target_var = torch.autograd.Variable(target.cuda(async=True))
output = model(input_var)
score_map = output[-1].data.cpu()
if flip:
flip_input_var = torch.autograd.Variable(
torch.from_numpy(shufflelr(inputs.clone().numpy())).float().cuda())
flip_output_var = model(flip_input_var)
flip_output = flip_back(flip_output_var[-1].data.cpu())
score_map += flip_output
# intermediate supervision
loss = 0
for o in output:
loss += criterion(o, target_var)
acc, batch_dists = accuracy(score_map, target.cpu(), idx, thr=0.07)
all_dists[:, i * args.val_batch:(i + 1) * args.val_batch] = batch_dists
preds = final_preds(score_map, meta['center'], meta['scale'], [64, 64])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
if debug:
gt_batch_img = batch_with_heatmap(inputs, target)
pred_batch_img = batch_with_heatmap(inputs, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
losses.update(loss.data[0], inputs.size(0))
acces.update(acc[0], inputs.size(0))
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
mean_error = torch.mean(all_dists)
auc = calc_metrics(all_dists) # this is auc of predicted maps and target.
print("=> Mean Error: {:.2f}, [email protected]: {} based on maps".format(mean_error*100., auc))
return losses.avg, acces.avg, predictions, auc
def train(loader, model, criterion, optimizer, netType, debug=False, flip=False):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
model.train()
end = time.time()
# rnn = torch.nn.LSTM(10, 20, 2)
# hidden = torch.autograd.Variable(torch.zeros((args.train_batch)))
gt_win, pred_win = None, None
bar = Bar('Training', max=len(loader))
for i, (inputs, target) in enumerate(loader):
data_time.update(time.time() - end)
input_var = torch.autograd.Variable(inputs.cuda())
target_var = torch.autograd.Variable(target.cuda(async=True))
if debug:
gt_batch_img = batch_with_heatmap(inputs, target)
# pred_batch_img = batch_with_heatmap(inputs, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
# plt.subplot(122)
# pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
# pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
output = model(input_var)
score_map = output[-1].data.cpu()
if flip:
flip_input_var = torch.autograd.Variable(
torch.from_numpy(shufflelr(inputs.clone().numpy())).float().cuda())
flip_output_var = model(flip_input_var)
flip_output = flip_back(flip_output_var[-1].data.cpu())
score_map += flip_output
# intermediate supervision
loss = 0
for o in output:
loss += criterion(o, target_var)
acc, _ = accuracy(score_map, target.cpu(), idx, thr=0.07)
losses.update(loss.data[0], inputs.size(0))
acces.update(acc[0], inputs.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg
