(self, class_num, no_local= 'on', gm_pool = 'on', arch='resnet50')
| 144 | |
| 145 | class embed_net(nn.Module): |
| 146 | def __init__(self, class_num, no_local= 'on', gm_pool = 'on', arch='resnet50'): |
| 147 | super(embed_net, self).__init__() |
| 148 | |
| 149 | self.thermal_module = thermal_module(arch=arch) |
| 150 | self.visible_module = visible_module(arch=arch) |
| 151 | self.base_resnet = base_resnet(arch=arch) |
| 152 | self.non_local = no_local |
| 153 | if self.non_local =='on': |
| 154 | layers=[3, 4, 6, 3] |
| 155 | non_layers=[0,2,3,0] |
| 156 | self.NL_1 = nn.ModuleList( |
| 157 | [Non_local(256) for i in range(non_layers[0])]) |
| 158 | self.NL_1_idx = sorted([layers[0] - (i + 1) for i in range(non_layers[0])]) |
| 159 | self.NL_2 = nn.ModuleList( |
| 160 | [Non_local(512) for i in range(non_layers[1])]) |
| 161 | self.NL_2_idx = sorted([layers[1] - (i + 1) for i in range(non_layers[1])]) |
| 162 | self.NL_3 = nn.ModuleList( |
| 163 | [Non_local(1024) for i in range(non_layers[2])]) |
| 164 | self.NL_3_idx = sorted([layers[2] - (i + 1) for i in range(non_layers[2])]) |
| 165 | self.NL_4 = nn.ModuleList( |
| 166 | [Non_local(2048) for i in range(non_layers[3])]) |
| 167 | self.NL_4_idx = sorted([layers[3] - (i + 1) for i in range(non_layers[3])]) |
| 168 | |
| 169 | |
| 170 | pool_dim = 2048 |
| 171 | self.l2norm = Normalize(2) |
| 172 | self.bottleneck = nn.BatchNorm1d(pool_dim) |
| 173 | self.bottleneck.bias.requires_grad_(False) # no shift |
| 174 | self.bottleneck_2 = nn.BatchNorm1d(pool_dim) |
| 175 | self.bottleneck_2.bias.requires_grad_(False) |
| 176 | |
| 177 | self.classifier = nn.Linear(pool_dim, class_num, bias=False) |
| 178 | |
| 179 | self.bottleneck.apply(weights_init_kaiming) |
| 180 | self.bottleneck_2.apply(weights_init_kaiming) |
| 181 | self.classifier.apply(weights_init_classifier) |
| 182 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) |
| 183 | self.gm_pool = gm_pool |
| 184 | |
| 185 | def forward(self, x1, x2, modal=0): |
| 186 | if modal == 0: |
nothing calls this directly
no test coverage detected