def evaluate_model(trained_model, data_loader): net = CrowdCounter() network.load_net(trained_model, net) net.cuda() net.eval() mae = 0.0 mse = 0.0 for blob in data_loader: im_data = blob['data'] gt_data = blob['gt_density'] density_map = net.forward(im_data, gt_data) density_map = density_map.data.cpu().numpy() gt_count = np.sum(gt_data) et_count = np.sum(density_map) mae += abs(gt_count - et_count) mse += ((gt_count - et_count) * (gt_count - et_count)) mae = mae / data_loader.get_num_samples() mse = np.sqrt(mse / data_loader.get_num_samples()) return mae, mse
re_cnt = False t = Timer() t.tic() data_loader = ImageDataLoader(train_path, train_gt_path, shuffle=True, gt_downsample=True, pre_load=True) data_loader_val = ImageDataLoader(val_path, val_gt_path, shuffle=False, gt_downsample=True, pre_load=True) best_mae = sys.maxsize for epoch in range(start_step, end_step+1): step = -1 train_loss = 0 for blob in data_loader: step = step + 1 im_data = blob['data'] gt_data = blob['gt_density'] density_map = net.forward(im_data, gt_data) loss = net.loss train_loss += loss.item() step_cnt += 1 optimizer.zero_grad() loss.backward() optimizer.step() if step % disp_interval == 0: duration = t.toc(average=False) fps = step_cnt / duration gt_count = np.sum(gt_data) density_map = density_map.data.cpu().numpy() et_count = np.sum(density_map) utils.save_results(im_data,gt_data,density_map, output_dir) log_text = 'epoch: %4d, step %4d, Time: %.4fs, gt_cnt: %4.1f, et_cnt: %4.1f' % (epoch,