CSRNet-pytorch版本代码复现

一、概述

在之前的博客中我有详细介绍过人群计数领域中密度图的生成方法，还有一篇CSRNet的论文学习笔记。
人群计数之生成密度图
论文学习笔记：CSRNet: Dilated Convolutional Neural Networks for Understanding the Highly Congested Scenes
接下来我使用ShanghaiTech数据集对论文中提出的CSRNet模型进行了复现，其中关于数据增强等一些细节进行了省略，以方便初学者更快地了解人群计数相关的方法。

二、数据集预处理

本文使用ShanghaiTech part_B_final部分数据集对进行复现，由于没有对数据进行预处理和数据增强等操作，而part_A_final部分的图片尺寸大小不一样，所以在进行训练的时候只能设置批处理大小为1（batch_size=1）。
关于数据预处理部分的代码请参考：人群计数之生成密度图

运行代码之后我们可以看到在原来存放标注文件的文件夹生成每张图片对应的ground_truth密度图（h5文件）。

经过CSRNet网络模型之后输入图片的尺寸变为原来的1/8，所以生成的真实密度图大小为原图片尺寸的1/8，如果做了数据增强处理（例如裁剪、resize（）操作等），需要重新生成一次真实密度图。

三、构建CSRNet网络

CSRnet网络模型主要分为前端和后端网络，前端网络采用的是剔除了全连接层的VGG-16，后端网络采用了空洞卷积的操作，论文中为了对比不同模型的效果，也是搭建了四种空洞率不同的后端网络。网络结构图如下所示：

关于VGG网络的详细介绍可以参考博客：使用pytorch搭建自己的网络之VGG

class CSRNet(nn.Module):
    def __init__(self, load_weights=True):
        super(CSRNet, self).__init__()
        self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
        self.backend_feat = [512, 512, 512,256,128,64]
        self.frontend = make_layers(self.frontend_feat)
        self.backend = make_layers(self.backend_feat,in_channels = 512,dilation = True)
        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)

        if load_weights:  #加载VGG16的预训练模型
            vgg16 = models.vgg16(pretrained = True)
            self._initialize_weights()  #初始化权重
            #加载torchvision中的预训练模型和参数后通过state_dict()方法提取参数
            for i in range(len(self.frontend.state_dict().items())):
                list(self.frontend.state_dict().items())[i][1].data[:] = list(vgg16.state_dict().items())[i][1].data[:]

    def forward(self,x):
        x = self.frontend(x)
        x = self.backend(x)
        x = self.output_layer(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

这里为了方便起见，直接使用了预训练的VGG-16模型和参数。在初始化函数__init__中，首先调用_initialize_weights（）函数对所有层的参数进行初始化，然后使用VGG16预训练好的参数对其中属于VGG的层初始化。

最后补上make_layers（）函数：

def make_layers(cfg, in_channels = 3, batch_norm=False, dilation = False):
    if dilation:
        d_rate = 2
    else:
        d_rate = 1
    layers = []
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate,dilation = d_rate)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

论文中经过对比发现CSRNet-B的效果最好，所以后面的研究中均使用了结构B。这里要修改的话也很简单，注意修改d_rate 的值就好。

四、训练和测试

4.1 主函数

def main():
    # defining the location of dataset
    root = '../dataset/ShanghaiTech_Crowd_Counting_Dataset/'
    part_A_train = os.path.join(root, 'part_A_final/train_data', 'images')
    part_A_test = os.path.join(root, 'part_A_final/test_data', 'images')
    part_B_train = os.path.join(root, 'part_B_final/train_data', 'images')
    part_B_test = os.path.join(root, 'part_B_final/test_data', 'images')

    # 参数设置
    batch_size = 8
    num_epochs = 100
    lr = 1e-7
    momentum = 0.95
    decay = 5 * 1e-4

    #加载数据集
    train_iter,train_num = load_local_dataset(part_B_train,batch_size)
    print('train_num=',train_num)
    test_iter,test_num = load_local_dataset(part_B_test, batch_size)
    print('test_num=',test_num)

    #使用所有GPU进行训练
    net = torch.nn.DataParallel(model.CSRNet()).cuda()

    #定义损失函数和优化器
    loss = torch.nn.MSELoss(reduction='sum').cuda()
    optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=momentum, weight_decay=decay)

    for i in range(num_epochs):
        print('第%d次训练' % (i+1) )
        train(train_iter, net, loss, optimizer)
        torch.save(net.module.state_dict(), f'./checkpoint/CSRNet_{str(i+1).zfill(4)}.pt')#保存模型
        validate(test_iter, net,test_num)

首先定义相关参数，然后调用load_local_dataset（）函数加载数据集，最后进行训练和测试，并保存模型。由于博主有多块GPU，所以使用了多块GPU共同训练以增加训练速度。代码片段是torch.nn.DataParallel(model.CSRNet())，如果使用此行代码，在保存训练模型并调用模型的时候要注意，使用其他方法可能会报错。

4.2 加载数据集

# 首先继承上面的dataset类。然后在__init__()方法中得到图像的路径，将图像路径组成一个数组，这样在__getitim__()中就可以直接读取：
class MyDataset(Dataset):  # 创建自己的类：MyDataset,这个类是继承的torch.utils.data.Dataset
    def __init__(self, img_paths, transform=None, target_transform=None):  # 初始化一些需要传入的参数
        super(MyDataset, self).__init__()  # 对继承自父类的属性进行初始化
        self.img_paths = img_paths
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):  # 这个方法是必须要有的，用于按照索引读取每个元素的具体内容
        img_path = self.img_paths[index] #根据索引获取图片路径
        img = Image.open(img_path).convert('RGB')

        gt_path = img_path.replace('.jpg', '.h5').replace('images', 'ground_truth')
        gt_file = h5py.File(gt_path,'r')
        target = np.asarray(gt_file['density'])
        target = cv2.resize(target, (target.shape[1] // 8, target.shape[0] // 8), interpolation=cv2.INTER_CUBIC) * 64

        if self.transform is not None:
            img = self.transform(img)  # 数据标签转换为Tensor

        return img, target  # return回哪些内容，那么我们在训练时循环读取每个batch时，就能获得哪些内容

    def __len__(self):  # 这个函数也必须要写，它返回的是数据集的长度，也就是多少张图片，要和loader的长度作区分
        return len(self.img_paths)

继承dataset类并重写该函数，然后定义load_local_dataset（）函数来加载数据集，返回的是数据集的迭代器和数据集大小。

def load_local_dataset(path_sets, batch_size=8):
    img_paths = []
    path_sets=[path_sets]
    for path in path_sets:
        for img_path in glob.glob(os.path.join(path, '*.jpg')):
            img_paths.append(img_path)
    # 加载数据集
    datasets =MyDataset(img_paths, transform=transform)
    # 调用DataLoader来创建数据集的迭代器
    dataset_iter = DataLoader(dataset=datasets, batch_size=batch_size, shuffle=True)
    return dataset_iter,len(img_paths)

附上图像的初始化操作代码：

# 图像的初始化操作
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

4.3 训练

def train(train_iter, net, criterion, optimizer):
    #开始训练
    net.train()  # 启用 BatchNormalization 和 Dropout
    for (img, target) in train_iter:
        img = img.cuda()
        target = torch.unsqueeze(target, 1).cuda()
        output = net(img) #调用模型进行训练

        #损失函数
        loss = criterion(output, target)
        #反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

这里没有什么复杂的操作，唯一要注意的是target通过torch.unsqueeze操作进行维度扩充，使得target和output的维度相同，能够计算损失函数。

4.4 测试

def validate(test_iter, net, dataset_num):
    net.eval() #开始测试
    mae = 0
    mse = 0

    for (img, target) in test_iter:
        img = img.cuda()
        target = torch.unsqueeze(target, 1).cuda()
        output = net(img)

        mae += abs(output.data.sum() - target.data.sum())
        mse += pow((output.data.sum() - target.data.sum()),2)
    mae = mae / dataset_num
    mse = pow((mse / dataset_num),0.5)

    print('MAE {mae:.3f} MSE {mse:.3f} '.format(mae = mae,mse = mse))

为了追求代码的极致简单，这里也只计算了MAE和MSE。

4.5 补充

1、在训练代码中均使用了GPU进行训练，如果你想使用CPU训练的话，修改一下代码net = torch.nn.DataParallel(model.CSRNet()).cuda()并去掉所有.cuda()后缀即可。
2、此blog仅提供部分代码和相关讲解，如果想要获取全部资源可以通过下面链接：
https://download.csdn.net/download/qq_40356092/12811426
博主也会提供复现代码的所有指导。

五、测试单张图片

以下内容的代码请下载完整资源包：CSRNet-pytorch

原始图片：

真实密度图：

预测的密度图：

预测人数为： 19.395416
真实人数为： 21.089602

六、 Write in the end

由于在复现的时候参考了大量的代码，所以和其他博主/教程中提供的代码可能有相似之处，但是因为时间久远，所以很难找到相应的参考链接，在这里就不提供了~

如果你对以上内容有任何疑问，均可以在留言区进行留言评论，博主看到了一定会及时回复。