| 343 | |
| 344 | |
| 345 | class FAN_use(nn.Module): |
| 346 | def __init__(self): |
| 347 | super(FAN_use, self).__init__() |
| 348 | self.num_modules = 1 |
| 349 | |
| 350 | # Base part |
| 351 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) |
| 352 | self.bn1 = nn.BatchNorm2d(64) |
| 353 | self.conv2 = ConvBlock(64, 128) |
| 354 | self.conv3 = ConvBlock(128, 128) |
| 355 | self.conv4 = ConvBlock(128, 256) |
| 356 | |
| 357 | # Stacking part |
| 358 | hg_module = 0 |
| 359 | self.add_module("m" + str(hg_module), HourGlass(1, 4, 256)) |
| 360 | self.add_module("top_m_" + str(hg_module), ConvBlock(256, 256)) |
| 361 | self.add_module( |
| 362 | "conv_last" + str(hg_module), |
| 363 | nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0), |
| 364 | ) |
| 365 | self.add_module( |
| 366 | "l" + str(hg_module), nn.Conv2d(256, 68, kernel_size=1, stride=1, padding=0) |
| 367 | ) |
| 368 | self.add_module("bn_end" + str(hg_module), nn.BatchNorm2d(256)) |
| 369 | |
| 370 | if hg_module < self.num_modules - 1: |
| 371 | self.add_module( |
| 372 | "bl" + str(hg_module), |
| 373 | nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0), |
| 374 | ) |
| 375 | self.add_module( |
| 376 | "al" + str(hg_module), |
| 377 | nn.Conv2d(68, 256, kernel_size=1, stride=1, padding=0), |
| 378 | ) |
| 379 | |
| 380 | self.avgpool = nn.MaxPool2d((2, 2), 2) |
| 381 | self.conv6 = nn.Conv2d(68, 1, 3, 2, 1) |
| 382 | self.fc = nn.Linear(28 * 28, 512) |
| 383 | self.bn5 = nn.BatchNorm2d(68) |
| 384 | self.relu = nn.ReLU(True) |
| 385 | |
| 386 | def forward(self, x): |
| 387 | x = F.relu(self.bn1(self.conv1(x)), True) |
| 388 | x = F.max_pool2d(self.conv2(x), 2) |
| 389 | x = self.conv3(x) |
| 390 | x = self.conv4(x) |
| 391 | |
| 392 | previous = x |
| 393 | |
| 394 | i = 0 |
| 395 | hg = self._modules["m" + str(i)](previous) |
| 396 | |
| 397 | ll = hg |
| 398 | ll = self._modules["top_m_" + str(i)](ll) |
| 399 | |
| 400 | ll = self._modules["bn_end" + str(i)](self._modules["conv_last" + str(i)](ll)) |
| 401 | tmp_out = self._modules["l" + str(i)](F.relu(ll)) |
| 402 | |