| 203 | |
| 204 | |
| 205 | class FirstCell(nn.Module): |
| 206 | |
| 207 | def __init__(self, in_chs_left, out_chs_left, in_chs_right, out_chs_right, pad_type=''): |
| 208 | super(FirstCell, self).__init__() |
| 209 | self.conv_1x1 = ActConvBn(in_chs_right, out_chs_right, 1, stride=1) |
| 210 | |
| 211 | self.act = nn.ReLU() |
| 212 | self.path_1 = nn.Sequential() |
| 213 | self.path_1.add_module('avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)) |
| 214 | self.path_1.add_module('conv', nn.Conv2d(in_chs_left, out_chs_left, 1, stride=1, bias=False)) |
| 215 | |
| 216 | self.path_2 = nn.Sequential() |
| 217 | self.path_2.add_module('pad', nn.ZeroPad2d((-1, 1, -1, 1))) |
| 218 | self.path_2.add_module('avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)) |
| 219 | self.path_2.add_module('conv', nn.Conv2d(in_chs_left, out_chs_left, 1, stride=1, bias=False)) |
| 220 | |
| 221 | self.final_path_bn = nn.BatchNorm2d(out_chs_left * 2, eps=0.001, momentum=0.1) |
| 222 | |
| 223 | self.comb_iter_0_left = BranchSeparables(out_chs_right, out_chs_right, 5, 1, pad_type) |
| 224 | self.comb_iter_0_right = BranchSeparables(out_chs_right, out_chs_right, 3, 1, pad_type) |
| 225 | |
| 226 | self.comb_iter_1_left = BranchSeparables(out_chs_right, out_chs_right, 5, 1, pad_type) |
| 227 | self.comb_iter_1_right = BranchSeparables(out_chs_right, out_chs_right, 3, 1, pad_type) |
| 228 | |
| 229 | self.comb_iter_2_left = create_pool2d('avg', 3, 1, count_include_pad=False, padding=pad_type) |
| 230 | |
| 231 | self.comb_iter_3_left = create_pool2d('avg', 3, 1, count_include_pad=False, padding=pad_type) |
| 232 | self.comb_iter_3_right = create_pool2d('avg', 3, 1, count_include_pad=False, padding=pad_type) |
| 233 | |
| 234 | self.comb_iter_4_left = BranchSeparables(out_chs_right, out_chs_right, 3, 1, pad_type) |
| 235 | |
| 236 | def forward(self, x, x_prev): |
| 237 | x_relu = self.act(x_prev) |
| 238 | x_path1 = self.path_1(x_relu) |
| 239 | x_path2 = self.path_2(x_relu) |
| 240 | x_left = self.final_path_bn(torch.cat([x_path1, x_path2], 1)) |
| 241 | x_right = self.conv_1x1(x) |
| 242 | |
| 243 | x_comb_iter_0_left = self.comb_iter_0_left(x_right) |
| 244 | x_comb_iter_0_right = self.comb_iter_0_right(x_left) |
| 245 | x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right |
| 246 | |
| 247 | x_comb_iter_1_left = self.comb_iter_1_left(x_left) |
| 248 | x_comb_iter_1_right = self.comb_iter_1_right(x_left) |
| 249 | x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right |
| 250 | |
| 251 | x_comb_iter_2_left = self.comb_iter_2_left(x_right) |
| 252 | x_comb_iter_2 = x_comb_iter_2_left + x_left |
| 253 | |
| 254 | x_comb_iter_3_left = self.comb_iter_3_left(x_left) |
| 255 | x_comb_iter_3_right = self.comb_iter_3_right(x_left) |
| 256 | x_comb_iter_3 = x_comb_iter_3_left + x_comb_iter_3_right |
| 257 | |
| 258 | x_comb_iter_4_left = self.comb_iter_4_left(x_right) |
| 259 | x_comb_iter_4 = x_comb_iter_4_left + x_right |
| 260 | |
| 261 | x_out = torch.cat([x_left, x_comb_iter_0, x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1) |
| 262 | return x_out |