手把手教你用ShanghaiTech数据集跑BL（Bayesian Loss）模型— 2.模型训练及测试

最新推荐文章于 2024-04-22 16:07:30 发布

johnhamtom

最新推荐文章于 2024-04-22 16:07:30 发布

阅读量735

点赞数 4

文章标签： python 深度学习计算机视觉图像处理人工智能

本文链接：https://blog.csdn.net/johnhamtom/article/details/129175550

版权

1.数据准备

经过上一篇博客（https://blog.csdn.net/johnhamtom/article/details/128731524?spm=1001.2014.3001.5502）的数据处理介绍，现在已经将数据处理为了如下格式，现在可以开始训练了。
在这里插入图片描述

2.模型训练

到了模型训练阶段就比较容易了，本文使用的操作系统是ubuntu系统，使用的软件是Pycharm + Anaconda，其中Pycharm 作为IDE来调试代码，Anaconda用来管理python模拟环境。
本文使用的是python3.8的环境，其它所需的库环境都安装了最新的，将上述环境准备好后，首先将utils目录下的regression_trainer.py的datasets.crowd改为datasets.crowd_sh，如下图所示。（crowd_sh.py是由官方提供的用于处理shanghaiTech数据集的代码）
在这里插入图片描述
然后打开train.py文件里设置数据文件夹目录、训练权重保存文件目录、参数等，其中–data-dir后面设置的是数据文件路径，引到数据根目录下即可，参考上述图片中显示的路径的话是写到part_A_final（训练PartA部分）或者part_B_final（训练PartB部分）；–save-dir存放权重保存路径；后面的还有learning rate、weight decay等参数的设置，在train.py里给的比较清楚了。
最后打开anaconda prompt 定位到项目文件夹目录位置，开启模拟环境并输入python train.py即可开始训练，如下图所示，里面的参数我使用的是官方给的默认参数，batch_size为1，训练轮次为1000轮，在600轮的时候开启验证，每五轮验证一次。
在这里插入图片描述

训练了约4小时左右就训练好了，ShanghaiTech数据集的PartA部分的场景比较复杂，验证时的MAE能到12+，而PartB部分的场景就没那么复杂，验证时的MAE在2左右，如下图所示。

Part_A Part_B

3.模型测试

模型测试阶段要做的事就用训练好的模型去预测测试集中的人群图片中的人群数量，本文将会使用Accuracy、MAE（Mean Average Error）、RMSE（Root Mean Square Error）指标去评估模型的预测结果。
Accuracy、MAE、RMSE的公式分别如下：
在这里插入图片描述

其中假设共有n张图片， diffi 指的是第 i 张图片的预测人数与真实人数的差的绝对值，GTi 指的是第i张图片的真实值。
接着调整下test.py文件的代码以达到需要，首先还是将它的datasets.crowd改为datasets.crowd_sh，跟模型训练阶段做的事一样，然后调整好数据集路径和权重读取路径，这个在代码中的位置也很清楚，接着新建两个dataframe在循环中存储每一张图片和所有图片的名字、GT值、预测值、准确率等信息，完整代码如下所示，详细信息看代码。

import torch
import os
import numpy as np
from datasets.crowd_sh import Crowd
from models.vgg import vgg19
import argparse
import pandas as pd
args = None


def parse_args():
    parser = argparse.ArgumentParser(description='Test ')
    parser.add_argument('--data-dir', default='./Shanghai/part_B_final',
                        help='training data directory')
    parser.add_argument('--save-dir', default='./result/people/Shanghai_B',
                        help='model directory')
    parser.add_argument('--device', default='0', help='assign device')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.device.strip()  # set vis gpu

    # 我这边把val修改为了test
    datasets = Crowd(os.path.join(args.data_dir, 'test'), 512, 8, is_gray=False, method='val')
    dataloader = torch.utils.data.DataLoader(datasets, 1, shuffle=False,
                                             num_workers=8, pin_memory=False)
    model = vgg19()
    device = torch.device('cuda')
    model.to(device)
    model.load_state_dict(torch.load(os.path.join(args.save_dir, 'best_model.pth'), device))
    epoch_minus = []
    # 两个dataframe在循环中存储每一张图片和所有图片的名字、GT值、预测值、准确率等信息
    df = pd.DataFrame(columns=['Name', 'GT', 'Pre','Accuracy'])
    dfNum = pd.DataFrame(columns=['Accuracy','AverageMAE','RMSE'])
    sum = 0 # 记录图片数量
    totalAcc = 0 # 记录总准确率
    for inputs, count, name in dataloader:
        inputs = inputs.to(device)
        assert inputs.size(0) == 1, 'the batch size should equal to 1'
        with torch.set_grad_enabled(False):
            outputs = model(inputs)
            temp_minu = count[0].item() - torch.sum(outputs).item()
            # 计算每一张图片的准确信息，并记录总准确率
            acc = 1-abs((torch.sum(outputs).item()) - count[0].item())/count[0].item()
            totalAcc = acc + totalAcc
            new = [name, count[0].item() , torch.sum(outputs).item() , acc]
            df.loc[sum] = new
            print(name, temp_minu, count[0].item(), torch.sum(outputs).item())
            epoch_minus.append(temp_minu)
            sum = sum + 1
            
    mse = np.sqrt(np.mean(np.square(epoch_minus)))
    mae = np.mean(np.abs(epoch_minus))
    newdfNum = [totalAcc/sum,mae,mse]
    dfNum.loc[0] = newdfNum
    epoch_minus = np.array(epoch_minus)
    log_str = 'Final Test: mae {}, mse {}'.format(mae, mse)
    print(log_str)
    filename = "shanghaiTech_Part_B"
    df.to_csv(r"./people/"+filename+"-Images.csv", index=False)  # CSV文件路径
    dfNum.to_csv(r"./people/"+filename+"-Total.csv", index=False)