Linhz/ViMNer

1import torch
2import torch.nn as nn
3from torch.autograd import Variable
4import torch.nn.functional as F
5
6
7class myResnet(nn.Module):
8    def __init__(self, resnet, if_fine_tune, device):
9        super(myResnet, self).__init__()
10        self.resnet = resnet
11        self.if_fine_tune = if_fine_tune
12        self.device = device
13
14    def forward(self, x, att_size=7):
15        # x shape batch_size * channels * 224 * 224
16        # 32 * 3 * 224 * 224
17
18        # batch_size * channels * 112 * 112
19        # 32 * 64 * 112 * 112
20        x = self.resnet.conv1(x)
21        x = self.resnet.bn1(x)
22        x = self.resnet.relu(x)
23
24        # 32 * 256 * 56 * 56
25        x = self.resnet.maxpool(x)
26
27        # 32 * 512 * 56 * 56
28        x = self.resnet.layer1(x)
29        # 32 * 512 * 28 * 28
30        x = self.resnet.layer2(x)
31        # 32 * 1024 * 14 * 14
32        x = self.resnet.layer3(x)
33        # 32 * 2048 * 7 * 7
34        x = self.resnet.layer4(x)
35
36        # 32 * 2048
37        fc = x.mean(3).mean(2)
38        # 32 * 2048 * 7 * 7
39        att = F.adaptive_avg_pool2d(x, [att_size, att_size])
40
41        # 32 * 2048 * 1 * 1
42        x = self.resnet.avgpool(x)
43        # 32 * 2048 * 2048
44        x = x.view(x.size(0), -1)
45
46        if not self.if_fine_tune:
47            x = Variable(x.data)
48            fc = Variable(fc.data)
49            att = Variable(att.data)
50
51        return x, fc, att