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()