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())
