人群计数CSRNet的pytorch实现

本文中对CSRNet: Dilated Convolutional Neural Networks for Understanding the Highly Congested Scenes（CVPR 2018）中的模型进行pytorch实现

import torch;import torch.nn as nn
from torchvision.models import vgg16
vgg=vgg16(pretrained=1)

import warnings
warnings.filterwarnings("ignore")
vgg10=torch.nn.Sequential(torch.nn.Conv2d(3,64,3,stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.Conv2d(64, 64, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.MaxPool2d(2,2),

                          torch.nn.Conv2d(64, 128, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.Conv2d(128, 128, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.MaxPool2d(2,2),

                          torch.nn.Conv2d(128, 256, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.Conv2d(256, 256, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.Conv2d(256, 256, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.MaxPool2d(2,2),  #尝试不进行下采样以达到不进行上采样

                          torch.nn.Conv2d(256, 512, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.Conv2d(512, 512, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                          torch.nn.Conv2d(512, 512, 3, stride=1,padding=1),
                          torch.nn.ReLU(inplace=True),
                         
)
class CSRNET(torch.nn.Module):
    def __init__(self, load_weights=False):
        super(CSRNET,self).__init__()
        self.vgg10=vgg10

       
        self.dconv1 = torch.nn.Conv2d(512, 512, 3, dilation=2, stride=1, padding=2)
        self.dconv2 = torch.nn.Conv2d(512, 512, 3, dilation=2, stride=1, padding=2)
        self.dconv3 = torch.nn.Conv2d(512, 512, 3, dilation=2, stride=1, padding=2)
        self.dconv4 = torch.nn.Conv2d(512, 256, 3, dilation=2, stride=1, padding=2)
        self.dconv5 = torch.nn.Conv2d(256, 128, 3, dilation=2, stride=1, padding=2)
        self.dconv6 = torch.nn.Conv2d(128, 64, 3, dilation=2, stride=1, padding=2)
        self.finalconv=torch.nn.Conv2d(64,1,1)
        self.relu=torch.nn.functional.relu
        if not load_weights:
            self.vgg10.load_state_dict(vgg.features[0:23].state_dict())
            


    def forward(self,x):
        y=self.vgg10(x)
        y = self.relu(self.dconv1(y))
        y = self.relu(self.dconv1(y))
        y = self.relu(self.dconv2(y))
        y = self.relu(self.dconv3(y))
        y = self.relu(self.dconv4(y))
        y = self.relu(self.dconv5(y))
        y = self.relu(self.dconv6(y))
        h=self.finalconv(y)