def val(epoch, opt, val_loader, model):
split = 'val'
model.eval()
# Initialize evaluations
wingloss = WingLoss(w=15,curvature=3)
cost = AverageMeter()
error2d = []
for i in range(2): error2d.append(AverageMeter())
num_iters = len(val_loader)
bar = Bar('==>', max=num_iters)
for i, (img, hmap, pts, norm_factor) in enumerate(val_loader):
img = img.to("cuda")
hmap = hmap.to("cuda")
pts = pts.to("cuda")
norm_factor = norm_factor.to("cuda")
# Constants
nb = img.shape[0]
nj = ref.num_landmarks
# Forward propagation
outputs = model(img)
# Compute cost
loss = 0.0
if opt.integral == 0:
loss += weighted_mse_loss(outputs[0], hmap)
elif opt.integral == 1:
loss += weighted_mse_loss(outputs[0], hmap)
loss += wingloss(outputs[1], pts)*opt.weight1
# Update evaluations
cost.update(loss.detach().item(), nb)
if opt.network == 'fan' or opt.network == 'fan-trained':
error2d[0].update(compute_error(outputs[opt.num_modules-1].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")))
else:
if opt.integral == 0:
error2d[0].update(compute_error(outputs[0].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")))
elif opt.integral == 1:
error2d[0].update(compute_error(outputs[0].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")))
error2d[1].update(compute_error_direct_efficientFAN((outputs[1]).detach(), pts.detach(), norm_factor, torch.ones([nb, nj]).to("cuda")))
msg = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:}'.format(epoch, i, num_iters, split=split, total=bar.elapsed_td, eta=bar.eta_td)
msg = '{} | Cost {cost.avg:.2f}'.format(msg, cost=cost)
for j in range(2):
msg = '{} | E{:d} {:.2f}'.format(msg, j, error2d[j].avg)
Bar.suffix = msg
bar.next()
bar.finish()
return cost.avg, [error2d[0].avg, error2d[1].avg]
def test(epoch, opt, test_loader, model, name):
split = 'test'
model.eval()
# Initialize evaluations
wingloss = WingLoss(w=15,curvature=3)
cost = AverageMeter()
error2d = []
for i in range(2): error2d.append(AverageMeter())
nme_list = []
time_all = 0.0
num_iters = len(test_loader)
bar = Bar('==>', max=num_iters)
# Save results
if opt.save_results == True:
heatmaps = []
landmarks = []
for i, (img, hmap, pts, norm_factor) in enumerate(test_loader):
img = img.to("cuda")
hmap = hmap.to("cuda")
pts = pts.to("cuda")
norm_factor = norm_factor.to("cuda")
# Constants
nb = img.shape[0]
nj = ref.num_landmarks
# Forward propagation
time1 = time.time()
outputs = model(img)
time2 = time.time()
time_all += (time2-time1)
# Compute cost
loss = 0.0
if opt.network == 'fan' or opt.network == 'fan-trained':
for j in range(opt.num_modules):
loss += weighted_mse_loss(outputs[j], hmap)
else:
if opt.integral == 0:
loss += weighted_mse_loss(outputs[0], hmap)
elif opt.integral == 1:
loss += weighted_mse_loss(outputs[0], hmap)
loss += wingloss(outputs[1], pts)*opt.weight1
# Update evaluations
cost.update(loss.detach().item(), nb)
if opt.network == 'fan' or opt.network == 'fan-trained':
error2d[0].update(compute_error(outputs[opt.num_modules-1].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")))
else:
if opt.integral == 0:
error2d[0].update(compute_error(outputs[0].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")))
elif opt.integral == 1:
error2d[0].update(compute_error(outputs[0].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")))
error2d[1].update(compute_error_direct_efficientFAN((outputs[1]).detach(), pts.detach(), norm_factor, torch.ones([nb, nj]).to("cuda")))
# Update NME
if opt.network == 'fan' or opt.network == 'fan-trained':
nme_list.append(compute_error(outputs[opt.num_modules-1].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")).item())
else:
if opt.integral == 0:
nme_list.append(compute_error(outputs[0].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")).item())
elif opt.integral == 1:
nme_list.append(compute_error_direct(outputs[1].detach(), pts.detach(), torch.ones([nb, nj]).to("cuda")).item())
# Save results
if opt.save_results == True:
heatmaps.append(outputs[0].detach().to("cpu"))
landmarks.append(outputs[1].detach().to("cpu"))
msg = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:}'.format(epoch, i, num_iters, split=split, total=bar.elapsed_td, eta=bar.eta_td)
msg = '{} | Cost {cost.avg:.2f}'.format(msg, cost=cost)
for j in range(2):
msg = '{} | E{:d} {:.2f}'.format(msg, j, error2d[j].avg)
Bar.suffix = msg
bar.next()
bar.finish()
# Final evaluation
nme_sorted = np.sort(nme_list)
num_samples = nme_sorted.shape[0]
bin = np.arange(0, 7.01, 0.05) #??
num_bin = bin.shape[0]
pdf = np.zeros(num_bin)
for i in range(num_samples):
val = nme_sorted[i]
idx = min(int(val/0.05)+1, num_bin-1)
pdf[idx] += 1
cdf = np.cumsum(pdf)
cdf = cdf / num_samples * 100.0
resultname = os.path.join(opt.save_dir, 'result_%s_cdf.txt' % name)
np.savetxt(resultname, cdf)
# Compute AUC
AUC = 0.0
for i in range(num_bin-1):
AUC += cdf[i+1]
AUC /= (num_bin-1)
print('AUC for 7%% = %.2f' % AUC)
resultname = os.path.join(opt.save_dir, 'result_%s_auc.txt' % name)
file = open(resultname, 'w')
file.write('AUC for 7%% = %.2f' % AUC)
file.close()
# Compute time
time_average = time_all / num_iters
fps = 1.0 / time_average
print('time for one frame = %.2f (s)' % time_average)
print('fps = %.2f (frame/s)' % fps)
resultname = os.path.join(opt.save_dir, 'result_%s_time.txt' % name)
file = open(resultname, 'w')
file.write('time for one frame = %.2f (ms)\n' % (time_average*1000.0))
file.write('fps = %.2f (frame/s)' % fps)
file.close()
# Save results
if opt.save_results == True:
heatmaps = torch.cat(heatmaps, 0)
filename = os.path.join(opt.save_dir, 'heatmaps_%s.pth' % name)
torch.save(heatmaps, filename)
landmarks = torch.cat(landmarks, 0)
filename = os.path.join(opt.save_dir, 'landmarks_%s.pth' % name)
torch.save(landmarks, filename)
if False:
# Constants for figure
fontsize = 12
figure_size = (11, 6)
# Figure
fig = plt.figure()
ax = plt.gca()
plt.plot(bin, cdf, linestyle='-', linewidth=2)
ax.set_xlabel('NME (%)')
ax.set_ylabel('Test Images (%)')
ax.set_title('300W-Testset, NME (%)', fontsize=fontsize)
ax.set_xlim([0., 7.])
ax.set_ylim([0., 100.])
ax.set_yticks(np.arange(0., 101., 10.))
plt.grid('on', linestyle='--', linewidth=0.5)
#fig.set_size_inches(np.asarray(figure_size))
plt.show()
resultname = os.path.join(opt.save_dir, 'result_cdf.jpg')
fig.savefig(resultname)
return cost.avg, [error2d[0].avg, error2d[1].avg]
def train(epoch, opt, train_loader, model, optimizer):
'''
해당 train함수는 1 epoch만 도는것임.
main에서 epoch만큼 train함수를 호출하기 때문에 총 epoch만큼 실행하게 된다.
- We set 𝜔 = 15 and 𝜀 = 3 for Wing loss [5], to calculate the 𝐿_P and the weight factor 𝜆 = e−5.
==> Wing loss, L_P가 뭔지 알아야함.
'''
split = 'train'
model.train()
# Wing loss test code
wingloss = WingLoss(w=15,curvature=3)
# Initialize evaluations
cost = AverageMeter() # cost는 loss이고
error2d = [] # error는 model의 성능을 나타내기 위한 지표. like NME, 정확도, mAP, ...
for i in range(2): error2d.append(AverageMeter())
num_iters = len(train_loader)
bar = Bar('==>', max=num_iters)
for i, (img, hmap, pts, norm_factor) in enumerate(train_loader):
img = img.to("cuda")
hmap = hmap.to("cuda")
pts = pts.to("cuda")
norm_factor = norm_factor.to("cuda")
# print(f'norm_factor shape: {norm_factor.shape}, norm type : {type(norm_factor)}')
# Constants
num_batch = img.shape[0] # batch size
num_joint = ref.num_landmarks # landmark number
# Forward propagation
outputs = model(img)
# Compute cost
loss = 0.0
# Compute loss
# 바꿔야하는 부분. mse_loss ==> wing loss
if opt.integral == 0:
loss += weighted_mse_loss(outputs[0], hmap)
elif opt.integral == 1:
loss += weighted_mse_loss(outputs[0], hmap)
loss += wingloss(outputs[1], pts)*opt.weight1 # wing loss test code
# Update model parameters with backpropagation
optimizer.zero_grad() #optimizer 선정도 opt에서 하는건가?
loss.backward()
optimizer.step()
# Update evaluations
cost.update(loss.detach().item(), num_batch) # detach()를 해야 tensor가 미분을 위해 추적하는 것을 막을 수 있다.
error2d[0].update(compute_error(outputs[0].detach(), pts.detach(), torch.ones([num_batch, num_joint]).to("cuda")))
error2d[1].update(compute_error_direct_efficientFAN((outputs[1]).detach(), pts.detach(), norm_factor, torch.ones([num_batch, num_joint]).to("cuda"))) #NME
# error2d[1].update(compute_error_direct((outputs[1]).detach(), pts.detach(), torch.ones([num_batch, num_joint]).to("cuda"))) #NME
# msg : epoch / 1 epoch 안에서의 iter / 1개 epoch 안에서의 총 iter / train,val / total : 1 epoch하면서 지금까지 걸린 시간 / eta : 예상 남은 시간 / cost / lr
msg = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:}'.format(epoch, i, num_iters, split=split, total=bar.elapsed_td, eta=bar.eta_td)
if split == 'train':
msg = '{} | LR: {:1.1e}'.format(msg, optimizer.param_groups[0]['lr'])
msg = '{} | Cost {cost.avg:.2f}'.format(msg, cost=cost)
for j in range(2):
msg = '{} | E{:d} {:.2f}'.format(msg, j, error2d[j].avg)
Bar.suffix = msg
bar.next()
bar.finish()
return cost.avg, [error2d[0].avg, error2d[1].avg]
def train_kd(epoch, opt, train_loader, model, teacher, optimizer):
'''
해당 train함수는 1 epoch만 도는것임.
main에서 epoch만큼 train함수를 호출하기 때문에 총 epoch만큼 실행하게 된다.
- We set 𝜔 = 15 and 𝜀 = 3 for Wing loss [5], to calculate the 𝐿_P and the weight factor 𝜆 = e−5.
==> Wing loss, L_P가 뭔지 알아야함.
'''
split = 'train'
model.train()
teacher.eval()
# Wing loss test code
wingloss = WingLoss(w=15,curvature=3)
# Initialize evaluations
cost = AverageMeter() # cost는 loss이고
kd_cost = AverageMeter()
error2d = [] # error는 model의 성능을 나타내기 위한 지표. like NME, 정확도, mAP, ...
for i in range(2): error2d.append(AverageMeter())
num_iters = len(train_loader)
bar = Bar('==>', max=num_iters)
for i, (img, hmap, pts, norm_factor) in enumerate(train_loader):
img = img.to("cuda")
hmap = hmap.to("cuda")
pts = pts.to("cuda")
norm_factor = norm_factor.to("cuda")
# print(f'norm_factor shape: {norm_factor.shape}, norm type : {type(norm_factor)}')
# Constants
num_batch = img.shape[0] # batch size
num_joint = ref.num_landmarks # landmark number
# Forward propagation
outputs = model(img)
# teacher로부터 나온 output에서는 gradient를 계산하지 않는다.
outputs_t = teacher(img)
for i in range(len(outputs_t)):
# outputs_t[2,3,4,5] : tensor type
outputs_t[i] = Variable(outputs_t[i], requires_grad=False)
# Compute loss
# 바꿔야하는 부분. mse_loss ==> wing loss
loss = 0.0
loss += weighted_mse_loss(outputs[0], hmap)
loss += wingloss(outputs[1], pts)*opt.weight1 # wing loss test code
fa_loss = 0.0
# ps_loss = 0.0
# feature alignment Loss
for i in range(2, 5+1):
assert (outputs[i].shape == outputs_t[i].shape)
fa_loss += feature_aligned_distill(outputs[i],outputs_t[i], weight=5e-4) # 5e-4 일때 gt loss : kd loss = 2:1
fa_loss = fa_loss*opt.weight2
# Patch Similarity Loss
'''for i in range(2, 5+1):
assert (len(outputs[i]) == len(outputs_t[i]))
ps_loss += patch_similarity_distill(outputs[i],outputs_t[i], weight=200.0)'''
# print(f' ')
# print(f'fa_loss : {fa_loss}')
# print(f'ps_loss : {ps_loss}')
# print(f'loss : {loss}')
loss = loss + fa_loss
# Update model parameters with backpropagation
optimizer.zero_grad() #optimizer 선정도 opt에서 하는건가?
loss.backward()
optimizer.step()
# Update evaluations
cost.update((loss-fa_loss).detach().item(), num_batch) # detach()를 해야 tensor가 미분을 위해 추적하는 것을 막을 수 있다.
kd_cost.update(fa_loss.detach().item(), num_batch)
error2d[0].update(compute_error(outputs[0].detach(), pts.detach(), torch.ones([num_batch, num_joint]).to("cuda")))
error2d[1].update(compute_error_direct_efficientFAN((outputs[1]).detach(), pts.detach(), norm_factor, torch.ones([num_batch, num_joint]).to("cuda"))) #NME
# msg : epoch / 1 epoch 안에서의 iter / 1개 epoch 안에서의 총 iter / train,val / total : 1 epoch하면서 지금까지 걸린 시간 / eta : 예상 남은 시간 / cost / lr
msg = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:}'.format(epoch, i, num_iters, split=split, total=bar.elapsed_td, eta=bar.eta_td)
if split == 'train':
msg = '{} | LR: {:1.1e}'.format(msg, optimizer.param_groups[0]['lr'])
msg = '{} | GT_Cost {cost.avg:.2f}'.format(msg, cost=cost)
msg = '{} | KD_Cost {cost.avg:.2f}'.format(msg, cost=kd_cost)
for j in range(2):
msg = '{} | E{:d} {:.2f}'.format(msg, j, error2d[j].avg)
Bar.suffix = msg
bar.next()
bar.finish()
return [cost.avg, kd_cost.avg], [error2d[0].avg, error2d[1].avg]