class MiniMobileNet(d2l.Classifier):
def __init__(self, arch=((64, 1), (128, 2), (128, 1), (256, 2), (256, 1),
(512, 2), (512, 1)),
lr=0.1, num_classes=10):
super().__init__()
self.save_hyperparameters()
layers = [nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1,
bias=False),
nn.BatchNorm2d(32), nn.ReLU()]
c = 32
for c_out, stride in arch:
layers.append(dws_block(c, c_out, stride))
c = c_out
layers += [nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(),
nn.Linear(c, num_classes)]
self.net = nn.Sequential(*layers)
def configure_optimizers(self):
return torch.optim.SGD(self.parameters(), lr=self.lr, momentum=0.9)