1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45
| class ResNet18(nn.Module): def __init__(self, BasicBlock, num_classes=10) -> None: super(ResNet18, self).__init__() self.in_channels = 64 self.conv1 = nn.Sequential( nn.Conv2d(3,64,kernel_size=7,stride=2,padding=3,bias=False), nn.BatchNorm2d(64), nn.MaxPool2d(kernel_size=3, stride=2, padding=1) ) self.conv2 = self._make_layer(BasicBlock,64,[[1,1],[1,1]])
self.conv3 = self._make_layer(BasicBlock,128,[[2,1],[1,1]])
self.conv4 = self._make_layer(BasicBlock,256,[[2,1],[1,1]])
self.conv5 = self._make_layer(BasicBlock,512,[[2,1],[1,1]]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512, num_classes)
def _make_layer(self, block, out_channels, strides): layers = [] for stride in strides: layers.append(block(self.in_channels, out_channels, stride)) self.in_channels = out_channels return nn.Sequential(*layers) def forward(self, x): out = self.conv1(x) out = self.conv2(out) out = self.conv3(out) out = self.conv4(out) out = self.conv5(out)
out = self.avgpool(out) out = out.reshape(x.shape[0], -1) out = self.fc(out) return out
|