loss = LOSS(den,es_den)
optimizer.zero_grad()
loss.backward()
optimizer.step()
trainloss.update(loss.item(), img.shape[0])
es_count = np.sum(es_den[0][0].cpu().detach().numpy())
gt_count = np.sum(den[0][0].cpu().detach().numpy())
escount.append(es_count)
gtcount.append(gt_count)
durantion = time.time()-trainstart
trainfps = step/durantion
trainmae,trainmse = eva_model(escount,gtcount)
writer.add_scalars('data/trainstate', {
'trainmse': trainmse,
'trainmae': trainmae}, epoch)
info = 'trianloss:{%.6f} @ trainmae:{%.3f} @ trainmse:{%.3f} @ fps:{%.3f}'%(trainloss.avg*10000,trainmae, trainmse,trainfps)
logger.info(info)
del escount[:]
del gtcount[:]
with torch.no_grad():
net.eval()
time_stamp = 0.0
for index,(timg,tden) in tqdm(enumerate(val_loader)):
start = time.time()
gt_count = np.sum(den[0][0].cpu().detach().numpy())
diff = abs(es_count-gt_count)
stop = time.time()
logger.info(name[0].split('/')[-1])
logger.info(' @ gt: %.3f vs %.3f:es @diff: %.3f'%(gt_count,es_count,diff))
duration+=stop-start
recod[name] = diff
escount.append(es_count)
gtcount.append(gt_count)
plt.subplot(131)
plt.title('raw image')
plt.imshow(img[0][0].cpu().detach().numpy())
plt.subplot(132)
plt.title('gtcount:%.2f' % gt_count)
plt.imshow(den[0][0].cpu().detach().numpy())
plt.subplot(133)
plt.title('escount:%.2f' % es_count)
plt.imshow(es_den[0][0].cpu().detach().numpy())
plt.savefig(saveimg + '.jpg'.format(name))
list1 = sorted(recod.items(), key=lambda x: x[1])
fps = duration/len(test_data)
MAE,MSE = eva_model(escount,gtcount)
logger.info('$$$$ MAE: %.3f - MSE: %.3f - FPS: %.3f $$$$'%(MAE,MSE,fps))
logger.info('!!!! sort start !!!!')
logger.info(list1)
logger.info('!!!! sort done !!!!')