A ResNet class that is similar to torchvision's but contains the following changes: - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool. - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride
| 91 | |
| 92 | |
| 93 | class ModifiedResNet(nn.Module): |
| 94 | """ |
| 95 | A ResNet class that is similar to torchvision's but contains the following changes: |
| 96 | - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool. |
| 97 | - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1 |
| 98 | - The final pooling layer is a QKV attention instead of an average pool |
| 99 | """ |
| 100 | |
| 101 | def __init__(self, layers, output_dim, heads, input_resolution=224, width=64): |
| 102 | super().__init__() |
| 103 | self.output_dim = output_dim |
| 104 | self.input_resolution = input_resolution |
| 105 | |
| 106 | # the 3-layer stem |
| 107 | self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False) |
| 108 | self.bn1 = nn.BatchNorm2d(width // 2) |
| 109 | self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False) |
| 110 | self.bn2 = nn.BatchNorm2d(width // 2) |
| 111 | self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False) |
| 112 | self.bn3 = nn.BatchNorm2d(width) |
| 113 | self.avgpool = nn.AvgPool2d(2) |
| 114 | self.relu = nn.ReLU(inplace=True) |
| 115 | |
| 116 | # residual layers |
| 117 | self._inplanes = width # this is a *mutable* variable used during construction |
| 118 | self.layer1 = self._make_layer(width, layers[0]) |
| 119 | self.layer2 = self._make_layer(width * 2, layers[1], stride=2) |
| 120 | self.layer3 = self._make_layer(width * 4, layers[2], stride=2) |
| 121 | self.layer4 = self._make_layer(width * 8, layers[3], stride=2) |
| 122 | |
| 123 | embed_dim = width * 32 # the ResNet feature dimension |
| 124 | self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim) |
| 125 | |
| 126 | def _make_layer(self, planes, blocks, stride=1): |
| 127 | layers = [Bottleneck(self._inplanes, planes, stride)] |
| 128 | |
| 129 | self._inplanes = planes * Bottleneck.expansion |
| 130 | for _ in range(1, blocks): |
| 131 | layers.append(Bottleneck(self._inplanes, planes)) |
| 132 | |
| 133 | return nn.Sequential(*layers) |
| 134 | |
| 135 | def forward(self, x): |
| 136 | def stem(x): |
| 137 | for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]: |
| 138 | x = self.relu(bn(conv(x))) |
| 139 | x = self.avgpool(x) |
| 140 | return x |
| 141 | |
| 142 | x = x.type(self.conv1.weight.dtype) |
| 143 | x = stem(x) |
| 144 | x = self.layer1(x) |
| 145 | x = self.layer2(x) |
| 146 | x = self.layer3(x) |
| 147 | x = self.layer4(x) |
| 148 | x = self.attnpool(x) |
| 149 | |
| 150 | return x |