This is just the code of the main program.
import os
import random
import numpy as np
import torch
from torch import optim
from torch.autograd import Variable
from torch.nn import NLLLoss2d
from torch.optim.lr_scheduler import StepLR,LambdaLR
import torchvision.transforms as standard_transforms
import torchvision.utils as vutils
from tensorboardX import SummaryWriter
from models.CC import CrowdCounter
from config import cfg
from loading_data import loading_data
from misc.utils import *
from misc.timer import Timer
import pdb
exp_name = cfg.TRAIN.EXP_NAME
if not os.path.exists(cfg.TRAIN.EXP_PATH):
os.mkdir(cfg.TRAIN.EXP_PATH)
writer = SummaryWriter(cfg.TRAIN.EXP_PATH+ ‘/’ + exp_name)
log_txt = cfg.TRAIN.EXP_PATH + ‘/’ + exp_name + ‘/’ + exp_name + ‘.txt’
pil_to_tensor = standard_transforms.ToTensor()
train_record = {‘best_mae’: 1e20, ‘mse’:1e20,‘corr_loss’: 0, ‘corr_epoch’: -1, ‘best_model_name’: ‘’}
train_set, train_loader, val_set, val_loader, restore_transform = loading_data()
_t = {‘iter time’ : Timer(),‘train time’ : Timer(),‘val time’ : Timer()}
rand_seed = cfg.TRAIN.SEED
if rand_seed is not None:
np.random.seed(rand_seed)
torch.manual_seed(rand_seed)
torch.cuda.manual_seed(rand_seed)
def main():
cfg_file = open(’./config.py’,“r”)
cfg_lines = cfg_file.readlines()
with open(log_txt, ‘a’) as f:
f.write(’’.join(cfg_lines) + ‘\n\n\n\n’)
if len(cfg.TRAIN.GPU_ID)==1:
torch.cuda.set_device(cfg.TRAIN.GPU_ID[0])
torch.backends.cudnn.benchmark = True
net = CrowdCounter().cuda()
if cfg.TRAIN.PRE_GCC:
net.load_state_dict(torch.load(cfg.TRAIN.PRE_GCC_MODEL))
net.train()
optimizer = optim.Adam(net.parameters(), lr=cfg.TRAIN.LR, weight_decay=1e-4)
scheduler = StepLR(optimizer, step_size=cfg.TRAIN.NUM_EPOCH_LR_DECAY, gamma=cfg.TRAIN.LR_DECAY)
i_tb = 0
# validate(val_loader, net, -1, restore_transform)
for epoch in range(cfg.TRAIN.MAX_EPOCH):
if epoch > cfg.TRAIN.LR_DECAY_START:
scheduler.step()
# training
_t[‘train time’].tic()
i_tb = train(train_loader, net, optimizer, epoch, i_tb)
_t[‘train time’].toc(average=False)
print(‘train time: {:.2f}s’.format(_t[‘train time’].diff))
print(’=’*20)
# validation
if epoch%cfg.VAL.FREQ==0 or epoch>cfg.VAL.DENSE_START:
_t[‘val time’].tic()
validate(val_loader, net, epoch, restore_transform)
_t[‘val time’].toc(average=False)
print(‘val time: {:.2f}s’.format(_t[‘val time’].diff))
def train(train_loader, net, optimizer, epoch, i_tb):
for i, data in enumerate(train_loader, 0):
_t[‘iter time’].tic()
img, gt_map, gt_cnt = data
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
optimizer.zero_grad()
pred_map = net(img, gt_map)
loss = net.loss
loss.backward()
optimizer.step()
pred_map = pred_map/100.
if (i + 1) % cfg.TRAIN.PRINT_FREQ == 0:
i_tb = i_tb + 1
writer.add_scalar(‘train_loss’, loss.data[0], i_tb)
_t[‘iter time’].toc(average=False)
print(’[ep %d][it %d][loss %.4f][lr %.4f][%.2fs]’ %
(epoch + 1, i + 1, loss.data[0], optimizer.param_groups[0][‘lr’]*10000, _t[‘iter time’].diff))
print(’ [cnt: gt: %.1f pred: %.2f]’ % (gt_cnt[0], pred_map[0,:,:].sum().data[0]))
return i_tb
def validate(val_loader, net, epoch, restore):
net.eval()
print(’=’*50)
val_loss = []
mae = 0.0
mse = 0.0
for vi, data in enumerate(val_loader, 0):
img, gt_map, gt_count = data
img = Variable(img, volatile=True).cuda()
gt_map = Variable(gt_map, volatile=True).cuda()
gt_count = gt_count.numpy()
pred_map = net(img, gt_map)
val_loss.append(net.loss.data)
pred_map = pred_map/100.
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map/100.
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt_tmp = np.sum(pred_map[i_img])
gt_count_tmp = gt_count[i_img]
mae += abs(gt_count_tmp-pred_cnt_tmp)
mse += ((gt_count_tmp-pred_cnt_tmp)*(gt_count_tmp-pred_cnt_tmp))
x = []
if vi==0:
for idx, tensor in enumerate(zip(img.cpu().data, pred_map, gt_map)):
if idx>cfg.VIS.VISIBLE_NUM_IMGS:
break
pil_input = restore(tensor[0])
pil_output = torch.from_numpy(tensor[1]/(tensor[1].max()+1e-10)).repeat(3,1,1)
pil_label = torch.from_numpy(tensor[2]/(tensor[2].max()+1e-10)).repeat(3,1,1)
x.extend([pil_to_tensor(pil_input.convert(‘RGB’)), pil_label, pil_output])
x = torch.stack(x, 0)
x = vutils.make_grid(x, nrow=3, padding=5)
writer.add_image(exp_name + ‘epoch’ + str(epoch+1), (x.numpy()*255).astype(np.uint8))
mae = mae/val_set.get_num_samples()
mse = np.sqrt(mse/val_set.get_num_samples())
loss = np.mean(np.array(val_loss))
writer.add_scalar(‘val_loss’, loss, epoch + 1)
writer.add_scalar(‘mae’, mae, epoch + 1)
writer.add_scalar(‘mse’, mse, epoch + 1)
snapshot_name = ‘ep_%d_mae_%.1f_mse_%.1f’ % (epoch + 1, mae, mse)
if mae < train_record[‘best_mae’]:
train_record[‘best_mae’] = mae
train_record[‘mse’] = mse
train_record[‘corr_epoch’] = epoch + 1
train_record[‘corr_loss’] = loss
train_record[‘best_model_name’] = snapshot_name
with open(log_txt, ‘a’) as f:
f.write(snapshot_name + ‘\n’)
# save model
to_saved_weight = []
to_saved_weight = net.state_dict()
torch.save(to_saved_weight, os.path.join(cfg.TRAIN.EXP_PATH, exp_name, snapshot_name + ‘.pth’))
print(’=’*50)
print(exp_name)
print(’ ‘+ ‘-’*20)
print(’ [mae %.2f mse %.2f], [val loss %.4f]’ % (mae, mse, loss))
print(’ ‘+ ‘-’*20)
# pdb.set_trace()
print(’[best] [mae %.2f mse %.2f], [loss %.4f], [epoch %d]’ % (train_record[‘best_mae’], train_record[‘mse’], train_record[‘corr_loss’], train_record[‘corr_epoch’]))
print(’=’*50)
net.train()
if name == ‘main’:
main()