SCAR的pytorch实现

本文所实现的网络来源于SCAR:Spatial-/Channel-wise Attention Regression Networks for Crowd Counting（Neurocompting 2019）

import torch;from torchvision import models
from torchvision.models import vgg16
import warnings;from torch import nn
warnings.filterwarnings("ignore")
vgg16 = vgg16(pretrained=True)
def initialize_weights(models):
    for model in models:
        real_init_weights(model)
import warnings
warnings.filterwarnings("ignore")
def real_init_weights(m):

    if isinstance(m, list):
        for mini_m in m:
            real_init_weights(mini_m)
    else:
        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.Linear):
            m.weight.data.normal_(0.0, std=0.01)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)
        elif isinstance(m,nn.Module):
            for mini_m in m.children():
                real_init_weights(mini_m)
        else:
            print( m )
class SCAR(torch.nn.Module):
    def __init__(self,loadwieght=False):
        super(SCAR,self).__init__()
        self.vgg10=vgg10
        if loadwieght==False:
            mod = models.vgg16(pretrained=True)
            initialize_weights(self.modules())
            self.vgg10.load_state_dict(mod.features[0:23].state_dict())
        
        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.relu = torch.nn.functional.relu
        self.SAM=SAM()
        self.CAM=CAM()
        self.finalconv=torch.nn.Conv2d(128,1,1)
        self.upsample=torch.nn.functional.upsample
    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))

        
        y_sa=self.SAM(y)
        
        y_ca=self.CAM(y)
       

        y=torch.cat((y_ca,y_sa),dim=1)
    
        y=self.finalconv(y)
       
        y=self.upsample(y,scale_factor=8)#由于进行了三次池化 因此8倍上取样
        return y


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),
                          #torch.nn.MaxPool2d(2),
)

class SAM(torch.nn.Module):
    def __init__(self):
        super(SAM,self).__init__()# SAM不改变输入到SAM中的x的shape
        self.q=torch.nn.Conv2d(64,64,1)
        self.k = torch.nn.Conv2d(64, 64, 1)
        self.v=torch.nn.Conv2d(64, 64, 1)
   

        self.lamda=torch.nn.Conv2d(64,64,1)
        self.bn=torch.nn.BatchNorm2d(64)


    def forward(self,x):
       
        N, C, H, W = x.size()
        
        q=self.q(x).view((N,-1,H*W)).permute(0,2,1) # HW*C
        k=self.q(x).view((N,-1,H*W))
        v=self.v(x).view((N,-1,H*W))
        mid=torch.bmm(q,k)
      
        attention=torch.nn.functional.softmax(mid,dim=-1)# HW*HW
        
        y=torch.bmm(v,attention)
        y=y.view((N,C,H,W))
        y=self.lamda(y)+x
     
        return y

class CAM(torch.nn.Module):
    def __init__(self):
        super(CAM,self).__init__()
        self.conv1=torch.nn.Conv2d(64,64,1)
        self.conv2 = torch.nn.Conv2d(64, 64, 1)
        self.bn = torch.nn.BatchNorm2d(64)


    def forward(self,x):
     
        N, C, H, W = x.size()
        q=self.conv1(x).view(N,C,-1)# C*HW
   
        k=self.conv1(x).view(N,-1,C) # HW*C
      
        attention_pre=torch.bmm(q,k)# C*C
        
        attention=torch.nn.functional.softmax(attention_pre,dim=-1)
        v=x.view(N,C,-1)
        cl2=torch.bmm(attention,v).view((N,C,H,W))
        cfinal=self.conv2(cl2)+x
       

        return cfinal
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158