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)
class Trainer():
def __init__(self, dataloader, cfg_data, pwd, cfg):
self.cfg_data = cfg_data
self.data_mode = cfg.DATASET
self.exp_name = cfg.EXP_NAME
self.exp_path = cfg.EXP_PATH
self.pwd = pwd
self.cfg = cfg
self.net_name = cfg.NET
self.net = CrowdCounter(cfg.GPU_ID, self.net_name, DA=True).cuda()
self.num_parameters = sum(
[param.nelement() for param in self.net.parameters()])
print('num_parameters:', self.num_parameters)
self.optimizer = optim.Adam(self.net.CCN.parameters(),
lr=cfg.LR,
weight_decay=1e-4)
# self.optimizer = optim.SGD(self.net.parameters(), cfg.LR, momentum=0.95,weight_decay=5e-4)
self.scheduler = StepLR(self.optimizer,
step_size=cfg.NUM_EPOCH_LR_DECAY,
gamma=cfg.LR_DECAY)
self.train_record = {
'best_mae': 1e20,
'best_mse': 1e20,
'best_model_name': '_'
}
self.hparam = {
'lr': cfg.LR,
'n_epochs': cfg.MAX_EPOCH,
'number of parameters': self.num_parameters,
'dataset': cfg.DATASET
} # ,'finetuned':cfg.FINETUNE}
self.timer = {
'iter time': Timer(),
'train time': Timer(),
'val time': Timer()
}
self.epoch = 0
self.i_tb = 0
'''discriminator'''
if cfg.GAN == 'Vanilla':
self.bce_loss = torch.nn.BCELoss()
elif cfg.GAN == 'LS':
self.bce_loss = torch.nn.MSELoss()
if cfg.NET == 'Res50':
self.channel1, self.channel2 = 1024, 128
self.D = [
FCDiscriminator(self.channel1, self.bce_loss).cuda(),
FCDiscriminator(self.channel2, self.bce_loss).cuda()
]
self.D[0].apply(weights_init())
self.D[1].apply(weights_init())
self.dis = self.cfg.DIS
self.d_opt = [
optim.Adam(self.D[0].parameters(),
lr=self.cfg.D_LR,
betas=(0.9, 0.99)),
optim.Adam(self.D[1].parameters(),
lr=self.cfg.D_LR,
betas=(0.9, 0.99))
]
self.scheduler_D = [
StepLR(self.d_opt[0],
step_size=cfg.NUM_EPOCH_LR_DECAY,
gamma=cfg.LR_DECAY),
StepLR(self.d_opt[1],
step_size=cfg.NUM_EPOCH_LR_DECAY,
gamma=cfg.LR_DECAY)
]
'''loss and lambdas here'''
self.lambda_adv = [cfg.LAMBDA_ADV1, cfg.LAMBDA_ADV2]
if cfg.PRE_GCC:
print('===================Loaded Pretrained GCC================')
weight = torch.load(cfg.PRE_GCC_MODEL)['net']
# weight=torch.load(cfg.PRE_GCC_MODEL)
try:
self.net.load_state_dict(convert_state_dict_gcc(weight))
except:
self.net.load_state_dict(weight)
# self.net=torch.nn.DataParallel(self.net, device_ids=cfg.GPU_ID).cuda()
'''modify dataloader'''
self.source_loader, self.target_loader, self.test_loader, self.restore_transform = dataloader(
)
self.source_len = len(self.source_loader.dataset)
self.target_len = len(self.target_loader.dataset)
print("source:", self.source_len)
print("target:", self.target_len)
self.source_loader_iter = cycle(self.source_loader)
self.target_loader_iter = cycle(self.target_loader)
if cfg.RESUME:
print('===================Loaded model to resume================')
latest_state = torch.load(cfg.RESUME_PATH)
self.net.load_state_dict(latest_state['net'])
self.optimizer.load_state_dict(latest_state['optimizer'])
self.scheduler.load_state_dict(latest_state['scheduler'])
self.epoch = latest_state['epoch'] + 1
self.i_tb = latest_state['i_tb']
self.train_record = latest_state['train_record']
self.exp_path = latest_state['exp_path']
self.exp_name = latest_state['exp_name']
self.writer, self.log_txt = logger(self.exp_path,
self.exp_name,
self.pwd,
'exp',
self.source_loader,
self.test_loader,
resume=cfg.RESUME,
cfg=cfg)
def forward(self):
print('forward!!')
# self.validate_V3()
with open(self.log_txt, 'a') as f:
f.write(str(self.net) + '\n')
f.write('num_parameters:' + str(self.num_parameters) + '\n')
for epoch in range(self.epoch, self.cfg.MAX_EPOCH):
self.epoch = epoch
# training
self.timer['train time'].tic()
self.train()
self.timer['train time'].toc(average=False)
if epoch > self.cfg.LR_DECAY_START:
self.scheduler.step()
self.scheduler_D[0].step()
self.scheduler_D[1].step()
print('train time: {:.2f}s'.format(self.timer['train time'].diff))
print('=' * 20)
self.net.eval()
# validation
if epoch % self.cfg.VAL_FREQ == 0 or epoch > self.cfg.VAL_DENSE_START:
self.timer['val time'].tic()
if self.data_mode in ['SHHA', 'SHHB', 'QNRF', 'UCF50', 'Mall']:
self.validate_V1()
elif self.data_mode is 'WE':
self.validate_V2()
elif self.data_mode is 'GCC':
self.validate_V3()
elif self.data_mode is 'NTU':
self.validate_V4()
# self.validate_train()
self.timer['val time'].toc(average=False)
print('val time: {:.2f}s'.format(self.timer['val time'].diff))
def train(self): # training for all datasets
self.net.train()
for i in range(max(len(self.source_loader), len(self.target_loader))):
torch.cuda.empty_cache()
self.timer['iter time'].tic()
img, gt_img = self.source_loader_iter.__next__()
tar, gt_tar = self.target_loader_iter.__next__()
img = Variable(img).cuda()
gt_img = Variable(gt_img).cuda()
tar = Variable(tar).cuda()
gt_tar = Variable(gt_tar).cuda()
#gen loss
# loss, loss_adv, pred, pred1, pred2, pred_tar, pred_tar1, pred_tar2 = self.gen_update(img,tar,gt_img,gt_tar)
self.optimizer.zero_grad()
for param in self.D[0].parameters():
param.requires_grad = False
for param in self.D[1].parameters():
param.requires_grad = False
# source
pred = self.net(img, gt_img)
loss = self.net.loss
if not self.cfg.LOSS_TOG:
loss.backward()
# target
pred_tar = self.net(tar, gt_tar)
loss_adv = self.D[self.dis].cal_loss(pred_tar[self.dis],
0) * self.lambda_adv[self.dis]
if not self.cfg.LOSS_TOG:
loss_adv.backward()
else:
loss += loss_adv
loss.backward()
#dis loss
loss_d = self.dis_update(pred, pred_tar)
self.d_opt[0].step()
self.d_opt[1].step()
self.optimizer.step()
if (i + 1) % self.cfg.PRINT_FREQ == 0:
self.i_tb += 1
self.writer.add_scalar('train_loss', loss.item(), self.i_tb)
self.writer.add_scalar('loss_adv', loss_adv.item(), self.i_tb)
self.writer.add_scalar('loss_d', loss_d.item(), self.i_tb)
self.timer['iter time'].toc(average=False)
print('[ep %d][it %d][loss %.4f][loss_adv %.8f][loss_d %.4f][lr %.8f][%.2fs]' % \
(self.epoch + 1, i + 1, loss.item(), loss_adv.item() if loss_adv else 0, loss_d.item(), self.optimizer.param_groups[0]['lr'],
self.timer['iter time'].diff))
print(' [cnt: gt: %.1f pred: %.2f]' %
(gt_img[0].sum().data / self.cfg_data.LOG_PARA,
pred[-1][0].sum().data / self.cfg_data.LOG_PARA))
print(' [tar: gt: %.1f pred: %.2f]' %
(gt_tar[0].sum().data / self.cfg_data.LOG_PARA,
pred_tar[-1][0].sum().data / self.cfg_data.LOG_PARA))
self.writer.add_scalar('lr', self.optimizer.param_groups[0]['lr'],
self.epoch + 1)
def gen_update(self, img, tar, gt_img, gt_tar):
pass
# return loss,loss_adv,pred,pred1,pred2,pred_tar,pred_tar1,pred_tar2
def dis_update(self, pred, pred_tar):
self.d_opt[self.dis].zero_grad()
for param in self.D[0].parameters():
param.requires_grad = True
for param in self.D[1].parameters():
param.requires_grad = True
#source
pred = [pred[0].detach(), pred[1].detach()]
loss_d = self.D[self.dis].cal_loss(pred[self.dis], 0)
if not self.cfg.LOSS_TOG:
loss_d.backward()
loss_D = loss_d
#target
pred_tar = [pred_tar[0].detach(), pred_tar[1].detach()]
loss_d = self.D[self.dis].cal_loss(pred_tar[self.dis], 1)
if not self.cfg.LOSS_TOG:
loss_d.backward()
loss_D += loss_d
if self.cfg.LOSS_TOG:
loss_D.backward()
return loss_D
def validate_train(self):
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
for img, gt_map in self.source_loader:
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
_, _, pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img]) / self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img]) / self.cfg_data.LOG_PARA
s_mae = abs(gt_count - pred_cnt)
s_mse = (gt_count - pred_cnt) * (gt_count - pred_cnt)
losses.update(self.net.loss.item())
maes.update(s_mae)
mses.update(s_mse)
loss = losses.avg
mae = maes.avg
mse = np.sqrt(mses.avg)
print("test on source domain")
print_NTU_summary(self.log_txt, self.epoch, [mae, mse, loss],
self.train_record)
def validate_V4(self): # validate_V4 for NTU
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
for vi, data in enumerate(self.test_loader, 0):
img, gt_map = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
_, _, pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img]) / self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img]) / self.cfg_data.LOG_PARA
s_mae = abs(gt_count - pred_cnt)
s_mse = (gt_count - pred_cnt) * (gt_count - pred_cnt)
losses.update(self.net.loss.item())
maes.update(s_mae)
mses.update(s_mse)
if vi == 0:
vis_results(self.exp_name, self.epoch, self.writer,
self.restore_transform, img, pred_map, gt_map)
loss = losses.avg
mae = maes.avg
mse = np.sqrt(mses.avg)
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mse', mse, self.epoch + 1)
self.train_record = update_model(self.net, self.optimizer,
self.scheduler, self.epoch, self.i_tb,
self.exp_path, self.exp_name,
[mae, mse, loss], self.train_record,
False, self.log_txt)
print_NTU_summary(self.log_txt, self.epoch, [mae, mse, loss],
self.train_record)
class Trainer():
def __init__(self, dataloader, cfg_data, pwd,cfg):
self.cfg_data = cfg_data
self.data_mode = cfg.DATASET
self.exp_name = cfg.EXP_NAME
self.exp_path = cfg.EXP_PATH
self.pwd = pwd
self.cfg=cfg
self.net_name = cfg.NET
self.net = CrowdCounter(cfg.GPU_ID,self.net_name).cuda()
self.num_parameters= sum([param.nelement() for param in self.net.parameters()])
print('num_parameters:',self.num_parameters)
self.optimizer = optim.Adam(self.net.CCN.parameters(), lr=cfg.LR, weight_decay=1e-4)
# self.optimizer = optim.SGD(self.net.parameters(), cfg.LR, momentum=0.95,weight_decay=5e-4)
self.scheduler = StepLR(self.optimizer, step_size=cfg.NUM_EPOCH_LR_DECAY, gamma=cfg.LR_DECAY)
self.train_record = {'best_mae': 1e20, 'best_mse':1e20, 'best_model_name': '_'}
self.hparam={'lr': cfg.LR, 'n_epochs': cfg.MAX_EPOCH,'number of parameters':self.num_parameters,'dataset':cfg.DATASET,'finetuned':cfg.FINETUNE}
self.timer = {'iter time' : Timer(),'train time' : Timer(),'val time' : Timer()}
self.epoch = 0
self.i_tb = 0
if cfg.PRE_GCC:
print('===================Loaded Pretrained GCC================')
weight=torch.load(cfg.PRE_GCC_MODEL)['net']
# weight=torch.load(cfg.PRE_GCC_MODEL)
try:
self.net.load_state_dict(convert_state_dict_gcc(weight))
except:
self.net.load_state_dict(weight)
# self.net=torch.nn.DataParallel(self.net, device_ids=cfg.GPU_ID).cuda()
self.train_loader, self.val_loader, self.restore_transform = dataloader()
if cfg.RESUME:
print('===================Loaded model to resume================')
latest_state = torch.load(cfg.RESUME_PATH)
self.net.load_state_dict(latest_state['net'])
self.optimizer.load_state_dict(latest_state['optimizer'])
self.scheduler.load_state_dict(latest_state['scheduler'])
self.epoch = latest_state['epoch'] + 1
self.i_tb = latest_state['i_tb']
self.train_record = latest_state['train_record']
self.exp_path = latest_state['exp_path']
self.exp_name = latest_state['exp_name']
self.writer, self.log_txt = logger(self.exp_path, self.exp_name, self.pwd, 'exp',self.train_loader, self.val_loader, resume=cfg.RESUME,cfg=cfg)
def forward(self):
# print('forward!!')
# self.validate_V3()
with open(self.log_txt, 'a') as f:
f.write(str(self.net) + '\n')
f.write('num_parameters:'+str(self.num_parameters)+'\n')
for epoch in range(self.epoch,self.cfg.MAX_EPOCH):
self.epoch = epoch
if epoch > self.cfg.LR_DECAY_START:
self.scheduler.step()
# training
self.timer['train time'].tic()
self.train()
self.timer['train time'].toc(average=False)
print( 'train time: {:.2f}s'.format(self.timer['train time'].diff) )
print( '='*20 )
self.net.eval()
# validation
if epoch%self.cfg.VAL_FREQ==0 or epoch>self.cfg.VAL_DENSE_START:
self.timer['val time'].tic()
if self.data_mode in ['SHHA', 'SHHB', 'QNRF', 'UCF50','Mall']:
self.validate_V1()
elif self.data_mode is 'WE':
self.validate_V2()
elif self.data_mode is 'GCC':
self.validate_V3()
elif self.data_mode is 'NTU':
self.validate_V4()
self.timer['val time'].toc(average=False)
print( 'val time: {:.2f}s'.format(self.timer['val time'].diff) )
def train(self): # training for all datasets
self.net.train()
for i, data in enumerate(self.train_loader, 0):
self.timer['iter time'].tic()
img, gt_map = data
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
self.optimizer.zero_grad()
pred_map = self.net(img, gt_map)
loss = self.net.loss
loss.backward()
self.optimizer.step()
if (i + 1) % self.cfg.PRINT_FREQ == 0:
self.i_tb += 1
self.writer.add_scalar('train_loss', loss.item(), self.i_tb)
self.timer['iter time'].toc(average=False)
print( '[ep %d][it %d][loss %.4f][lr %.6f][%.2fs]' % \
(self.epoch + 1, i + 1, loss.item(), self.optimizer.param_groups[0]['lr'], self.timer['iter time'].diff) )
print( ' [cnt: gt: %.1f pred: %.2f]' % (gt_map[0].sum().data/self.cfg_data.LOG_PARA, pred_map[0].sum().data/self.cfg_data.LOG_PARA) )
self.writer.add_scalar('lr', self.optimizer.param_groups[0]['lr'], self.epoch + 1)
def validate_V1(self):# validate_V1 for SHHA, SHHB, UCF-QNRF, UCF50
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
for vi, data in enumerate(self.val_loader, 0):
img, gt_map = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img])/self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img])/self.cfg_data.LOG_PARA
losses.update(self.net.loss.item())
maes.update(abs(gt_count-pred_cnt))
mses.update((gt_count-pred_cnt)*(gt_count-pred_cnt))
if vi==0:
vis_results(self.exp_name, self.epoch, self.writer, self.restore_transform, img, pred_map, gt_map)
mae = maes.avg
mse = np.sqrt(mses.avg)
loss = losses.avg
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mse', mse, self.epoch + 1)
self.train_record = update_model(self.net,self.optimizer,self.scheduler,self.epoch,self.i_tb,self.exp_path,self.exp_name, \
[mae, mse, loss],self.train_record,False,self.log_txt)
print_summary(self.log_txt,self.epoch,self.exp_name,[mae, mse, loss],self.train_record)
def validate_V2(self):# validate_V2 for WE
self.net.eval()
losses = AverageCategoryMeter(5)
maes = AverageCategoryMeter(5)
roi_mask = []
from datasets.WE.setting import cfg_data
from scipy import io as sio
for val_folder in cfg_data.VAL_FOLDER:
roi_mask.append(sio.loadmat(os.path.join(cfg_data.DATA_PATH,'test',val_folder + '_roi.mat'))['BW'])
for i_sub,i_loader in enumerate(self.val_loader,0):
mask = roi_mask[i_sub]
for vi, data in enumerate(i_loader, 0):
img, gt_map = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img,gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img])/self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img])/self.cfg_data.LOG_PARA
losses.update(self.net.loss.item(),i_sub)
maes.update(abs(gt_count-pred_cnt),i_sub)
if vi==0:
vis_results(self.exp_name, self.epoch, self.writer, self.restore_transform, img, pred_map, gt_map)
mae = np.average(maes.avg)
loss = np.average(losses.avg)
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mae_s1', maes.avg[0], self.epoch + 1)
self.writer.add_scalar('mae_s2', maes.avg[1], self.epoch + 1)
self.writer.add_scalar('mae_s3', maes.avg[2], self.epoch + 1)
self.writer.add_scalar('mae_s4', maes.avg[3], self.epoch + 1)
self.writer.add_scalar('mae_s5', maes.avg[4], self.epoch + 1)
self.train_record = update_model(self.net,self.optimizer,self.scheduler,self.epoch,self.i_tb,self.exp_path,self.exp_name, \
[mae, 0, loss],self.train_record,self.log_txt)
print_WE_summary(self.log_txt,self.epoch,[mae, 0, loss],self.train_record,maes)
# self.writer.add_hparams(self.hparam, {'best_mae': mae, 'best_mse':mse})
def validate_V3(self):# validate_V3 for GCC
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
c_maes = {'level':AverageCategoryMeter(9), 'time':AverageCategoryMeter(8),'weather':AverageCategoryMeter(7)}
c_mses = {'level':AverageCategoryMeter(9), 'time':AverageCategoryMeter(8),'weather':AverageCategoryMeter(7)}
for vi, data in enumerate(self.val_loader, 0):
img, gt_map, attributes_pt = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img])/self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img])/self.cfg_data.LOG_PARA
s_mae = abs(gt_count-pred_cnt)
s_mse = (gt_count-pred_cnt)*(gt_count-pred_cnt)
losses.update(self.net.loss.item())
maes.update(s_mae)
mses.update(s_mse)
attributes_pt = attributes_pt.squeeze()
c_maes['level'].update(s_mae,attributes_pt[i_img][0])
c_mses['level'].update(s_mse,attributes_pt[i_img][0])
c_maes['time'].update(s_mae,attributes_pt[i_img][1]/3)
c_mses['time'].update(s_mse,attributes_pt[i_img][1]/3)
c_maes['weather'].update(s_mae,attributes_pt[i_img][2])
c_mses['weather'].update(s_mse,attributes_pt[i_img][2])
if vi==0:
vis_results(self.exp_name, self.epoch, self.writer, self.restore_transform, img, pred_map, gt_map)
loss = losses.avg
mae = maes.avg
mse = np.sqrt(mses.avg)
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mse', mse, self.epoch + 1)
self.train_record = update_model(self.net,self.optimizer,self.scheduler,self.epoch,self.i_tb,self.exp_path,self.exp_name, \
[mae, mse, loss],self.train_record,False,self.log_txt)
print_GCC_summary(self.log_txt,self.epoch,[mae, mse, loss],self.train_record,c_maes,c_mses)
def validate_V4(self):# validate_V4 for NTU
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
for vi, data in enumerate(self.val_loader, 0):
img, gt_map = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img])/self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img])/self.cfg_data.LOG_PARA
s_mae = abs(gt_count-pred_cnt)
s_mse = (gt_count-pred_cnt)*(gt_count-pred_cnt)
losses.update(self.net.loss.item())
maes.update(s_mae)
mses.update(s_mse)
if vi==0:
vis_results(self.exp_name, self.epoch, self.writer, self.restore_transform, img, pred_map, gt_map)
loss = losses.avg
mae = maes.avg
mse = np.sqrt(mses.avg)
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mse', mse, self.epoch + 1)
self.train_record = update_model(self.net,self.optimizer,self.scheduler,self.epoch,self.i_tb,self.exp_path,self.exp_name, \
[mae, mse, loss],self.train_record,False,self.log_txt)
print_NTU_summary(self.log_txt,self.epoch,[mae, mse, loss],self.train_record)
# pruned_model_path = './exp/Res50_Original_GCC_Inducing_CAP_0.0001_epochs_100_Pruning/0.7/resnet50_GCC_pruned_0.7.pth.tar'
# pruned_model_path = './exp/VGG_Decoder_GCC_Pretrained_Pruning/0.4/VGG_Decoder_GCC_pruned_0.4.pth.tar'
# model_path='05-ResNet-50_all_ep_35_mae_32.4_mse_76.1.pth'
net = CrowdCounter(cfg.GPU_ID, cfg.NET)
# net = CrowdCounter(cfg.GPU_ID,cfg.NET,cfg=torch.load(pruned_model_path)['cfg'])
state_dict = torch.load(args.model_path)
try:
net.load_state_dict(state_dict['net'])
except KeyError:
net.load_state_dict(state_dict)
net.cuda()
net.eval()
sum([param.nelement() for param in net.parameters()])
def get_concat_h(im1, im2):
dst = Image.new('RGB', (im1.width + im2.width, im1.height))
dst.paste(im1, (0, 0))
dst.paste(im2, (im1.width, 0))
return dst
cm = plt.get_cmap('jet')
file_folder = []
file_name = []
'''
for file in glob.glob('/export/home/jinc0008/ntu_random_test/*'):
class Trainer():
def __init__(self, dataloader, cfg_data, pwd):
self.cfg_data = cfg_data
self.data_mode = cfg.DATASET
self.exp_name = cfg.EXP_NAME
self.exp_path = cfg.EXP_PATH
self.pwd = pwd
self.net_name = cfg.NET
self.net = CrowdCounter(cfg.GPU_ID, self.net_name).cuda()
self.optimizer = optim.Adam(self.net.parameters(),
lr=cfg.LR,
weight_decay=1e-4)
# self.optimizer = optim.SGD(self.net.parameters(), cfg.LR, momentum=0.95,weight_decay=5e-4)
self.scheduler = StepLR(self.optimizer,
step_size=cfg.NUM_EPOCH_LR_DECAY,
gamma=cfg.LR_DECAY)
self.train_record = {
'best_mae': 1e20,
'best_mse': 1e20,
'best_model_name': ''
}
self.timer = {
'iter time': Timer(),
'train time': Timer(),
'val time': Timer()
}
self.writer, self.log_txt = logger(self.exp_path, self.exp_name,
self.pwd, 'exp')
self.i_tb = 0
self.epoch = -1
if cfg.PRE_GCC:
self.net.load_state_dict(torch.load(cfg.PRE_GCC_MODEL))
self.train_loader, self.val_loader, self.restore_transform = dataloader(
)
def forward(self):
# self.validate_V1()
for epoch in range(cfg.MAX_EPOCH):
self.epoch = epoch
if epoch > cfg.LR_DECAY_START:
self.scheduler.step()
# training
self.timer['train time'].tic()
self.train()
self.timer['train time'].toc(average=False)
print 'train time: {:.2f}s'.format(self.timer['train time'].diff)
print '=' * 20
# validation
if epoch % cfg.VAL_FREQ == 0 or epoch > cfg.VAL_DENSE_START:
self.timer['val time'].tic()
if self.data_mode in ['SHHA', 'SHHB', 'QNRF', 'UCF50']:
self.validate_V1()
elif self.data_mode is 'WE':
self.validate_V2()
elif self.data_mode is 'GCC':
self.validate_V3()
self.timer['val time'].toc(average=False)
print 'val time: {:.2f}s'.format(self.timer['val time'].diff)
def train(self): # training for all datasets
self.net.train()
for i, data in enumerate(self.train_loader, 0):
self.timer['iter time'].tic()
img, gt_map = data
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
self.optimizer.zero_grad()
pred_map = self.net(img, gt_map)
loss = self.net.loss
loss.backward()
self.optimizer.step()
if (i + 1) % cfg.PRINT_FREQ == 0:
self.i_tb += 1
self.writer.add_scalar('train_loss', loss.item(), self.i_tb)
self.timer['iter time'].toc(average=False)
print '[ep %d][it %d][loss %.4f][lr %.4f][%.2fs]' % \
(self.epoch + 1, i + 1, loss.item(), self.optimizer.param_groups[0]['lr']*10000, self.timer['iter time'].diff)
print ' [cnt: gt: %.1f pred: %.2f]' % (
gt_map[0].sum().data / self.cfg_data.LOG_PARA,
pred_map[0].sum().data / self.cfg_data.LOG_PARA)
def validate_V1(self): # validate_V1 for SHHA, SHHB, UCF-QNRF, UCF50
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
time_sampe = 0
step = 0
for vi, data in enumerate(self.val_loader, 0):
img, gt_map = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img, gt_map)
step = step + 1
time_start1 = time.time()
test_map = self.net.test_forward(img)
time_end1 = time.time()
time_sampe = time_sampe + (time_end1 - time_start1)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
pred_cnt = np.sum(pred_map) / self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map) / self.cfg_data.LOG_PARA
losses.update(self.net.loss.item())
maes.update(abs(gt_count - pred_cnt))
mses.update((gt_count - pred_cnt) * (gt_count - pred_cnt))
if vi == 0:
vis_results(self.exp_name, self.epoch, self.writer,
self.restore_transform, img, pred_map, gt_map)
mae = maes.avg
mse = np.sqrt(mses.avg)
loss = losses.avg
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mse', mse, self.epoch + 1)
self.train_record = update_model(self.net, self.epoch, self.exp_path,
self.exp_name, [mae, mse, loss],
self.train_record, self.log_txt)
print_summary(self.exp_name, [mae, mse, loss], self.train_record)
print '\nForward Time: %fms' % (time_sampe * 1000 / step)
def validate_V2(self): # validate_V2 for WE
self.net.eval()
losses = AverageCategoryMeter(5)
maes = AverageCategoryMeter(5)
roi_mask = []
from datasets.WE.setting import cfg_data
from scipy import io as sio
for val_folder in cfg_data.VAL_FOLDER:
roi_mask.append(
sio.loadmat(
os.path.join(cfg_data.DATA_PATH, 'test',
val_folder + '_roi.mat'))['BW'])
for i_sub, i_loader in enumerate(self.val_loader, 0):
mask = roi_mask[i_sub]
for vi, data in enumerate(i_loader, 0):
img, gt_map = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(
pred_map[i_img]) / self.cfg_data.LOG_PARA
gt_count = np.sum(
gt_map[i_img]) / self.cfg_data.LOG_PARA
losses.update(self.net.loss.item(), i_sub)
maes.update(abs(gt_count - pred_cnt), i_sub)
if vi == 0:
vis_results(self.exp_name, self.epoch, self.writer,
self.restore_transform, img, pred_map,
gt_map)
mae = np.average(maes.avg)
loss = np.average(losses.avg)
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mae_s1', maes.avg[0], self.epoch + 1)
self.writer.add_scalar('mae_s2', maes.avg[1], self.epoch + 1)
self.writer.add_scalar('mae_s3', maes.avg[2], self.epoch + 1)
self.writer.add_scalar('mae_s4', maes.avg[3], self.epoch + 1)
self.writer.add_scalar('mae_s5', maes.avg[4], self.epoch + 1)
self.train_record = update_model(self.net, self.epoch, self.exp_path,
self.exp_name, [mae, 0, loss],
self.train_record, self.log_txt)
print_WE_summary(self.log_txt, self.epoch, [mae, 0, loss],
self.train_record, maes)
def validate_V3(self): # validate_V3 for GCC
self.net.eval()
losses = AverageMeter()
maes = AverageMeter()
mses = AverageMeter()
c_maes = {
'level': AverageCategoryMeter(9),
'time': AverageCategoryMeter(8),
'weather': AverageCategoryMeter(7)
}
c_mses = {
'level': AverageCategoryMeter(9),
'time': AverageCategoryMeter(8),
'weather': AverageCategoryMeter(7)
}
for vi, data in enumerate(self.val_loader, 0):
img, gt_map, attributes_pt = data
with torch.no_grad():
img = Variable(img).cuda()
gt_map = Variable(gt_map).cuda()
pred_map = self.net.forward(img, gt_map)
pred_map = pred_map.data.cpu().numpy()
gt_map = gt_map.data.cpu().numpy()
for i_img in range(pred_map.shape[0]):
pred_cnt = np.sum(pred_map[i_img]) / self.cfg_data.LOG_PARA
gt_count = np.sum(gt_map[i_img]) / self.cfg_data.LOG_PARA
s_mae = abs(gt_count - pred_cnt)
s_mse = (gt_count - pred_cnt) * (gt_count - pred_cnt)
losses.update(self.net.loss.item())
maes.update(s_mae)
mses.update(s_mse)
# attributes_pt = attributes_pt.squeeze()
# c_maes['level'].update(s_mae, attributes_pt[i_img][0])
# c_mses['level'].update(s_mse, attributes_pt[i_img][0])
# c_maes['time'].update(s_mae, attributes_pt[i_img][1] / 3)
# c_mses['time'].update(s_mse, attributes_pt[i_img][1] / 3)
# c_maes['weather'].update(s_mae, attributes_pt[i_img][2])
# c_mses['weather'].update(s_mse, attributes_pt[i_img][2])
# if vi == 0:
# vis_results(self.exp_name, self.epoch, self.writer, self.restore_transform, img, pred_map, gt_map)
loss = losses.avg
mae = maes.avg
mse = np.sqrt(mses.avg)
self.writer.add_scalar('val_loss', loss, self.epoch + 1)
self.writer.add_scalar('mae', mae, self.epoch + 1)
self.writer.add_scalar('mse', mse, self.epoch + 1)
self.train_record = update_model(self.net, self.epoch, self.exp_path,
self.exp_name, [mae, mse, loss],
self.train_record, self.log_txt)
print_GCC_summary(self.log_txt, self.epoch, [mae, mse, loss],
self.train_record, c_maes, c_mses)