def __init__(self, in_channels, n_classes, img_size): super(GoogleNet, self).__init__() self.head = nn.Sequential( ConvBlock(in_channels, 64, kernel_size=7, stride=2, padding=3), nn.MaxPool2d(kernel_size=3, stride=2, padding=1), ConvBlock(64, 192, kernel_size=3, stride=1, padding=1), nn.MaxPool2d(kernel_size=3, stride=2, padding=1)) if Params.isTrue('UseSqueezeExcitation'): InceptionBlock = SE_InceptionBlockv1 else: InceptionBlock = InceptionBlockv1 # Layers up to aux classifier 1 inc3a = InceptionBlock( InceptionParams(192, 64, 96, 128, 16, 32, 32, 256)) inc3b = InceptionBlock( InceptionParams(256, 128, 128, 192, 32, 96, 64, 480)) max1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) inc4a = InceptionBlock( InceptionParams(480, 192, 96, 208, 16, 48, 64, 512)) self.step1 = nn.Sequential(inc3a, inc3b, max1, inc4a) # Layers up to aux classifier 2 inc4b = InceptionBlock( InceptionParams(512, 160, 112, 224, 24, 64, 64, 512)) inc4c = InceptionBlock( InceptionParams(512, 128, 128, 256, 24, 64, 64, 512)) inc4d = InceptionBlock( InceptionParams(512, 112, 144, 288, 32, 64, 64, 528)) self.step2 = nn.Sequential(inc4b, inc4c, inc4d) # Layers up to main classifier inc4e = InceptionBlock( InceptionParams(528, 256, 160, 320, 32, 128, 128, 832)) max2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) inc5a = InceptionBlock( InceptionParams(832, 256, 160, 320, 32, 128, 128, 832)) inc5b = InceptionBlock( InceptionParams(832, 384, 192, 384, 48, 128, 128, 1024)) self.step3 = nn.Sequential(inc4e, max2, inc5a, inc5b) self.aux_classifier1 = nn.Sequential( nn.AvgPool2d(kernel_size=5, stride=3, padding=0), ConvBlock(512, 128, kernel_size=1, stride=1, padding=0), nn.Flatten(), nn.Linear(16 * 128, 1024), nn.ReLU(), nn.Dropout(0.7), nn.Linear(1024, n_classes)) self.aux_classifier2 = nn.Sequential( nn.AvgPool2d(kernel_size=5, stride=3, padding=0), ConvBlock(528, 128, kernel_size=1, stride=1, padding=0), nn.Flatten(), nn.Linear(16 * 128, 1024), nn.ReLU(), nn.Dropout(0.7), nn.Linear(1024, n_classes)) # Network is designed for 224 x 224 images. Here we introduce some flexibility # TODO: Aux Classifiers still assume 224x224 (although up to 256x256 should still work) test_img = torch.zeros((1, 3, img_size, img_size)) with torch.no_grad(): test_img = self.head(test_img) test_img = self.step1(test_img) test_img = self.step2(test_img) test_img = self.step3(test_img) out_img_size = test_img.shape[2] self.main_classifier = nn.Sequential( nn.AvgPool2d(kernel_size=out_img_size, stride=1, padding=0), nn.Dropout(0.4), nn.Flatten(), nn.Linear(1024, n_classes)) self.optimizer = Params.create_optimizer(self.parameters())
def __init__(self): self.net = vgg16(num_classes=Params.NumClasses) self.n_classes = Params.NumClasses self.optimizer = Params.create_optimizer(self.net.parameters())
def __init__(self): self.net = GoogleNet(Params.InChannels, Params.NumClasses, Params.CropSize) self.n_classes = Params.NumClasses self.optimizer = Params.create_optimizer(self.net.parameters())