def __init__(self):
self.load_datasets()
self.create_model()
if Params.isTrue('ContinueTrain'):
self.prepare_continued_training()
else:
self.start_epoch = 0
def training_step(self, images, targets):
self.optimizer.zero_grad()
result = self.net.forward(images)
result = self.wrap_result(result, True)
loss_record = self.calculate_losses(result, targets)
loss_record['final'].backward()
if Params.isTrue('ClipGradients'):
torch.nn.utils.clip_grad_norm_(self.net.parameters(),
Params.ClipGradientThreshold)
self.optimizer.step()
return loss_record
def prepare_logs(self):
os.makedirs( path.join(Params.CheckpointDir, Params.ModelName), exist_ok=True )
self.logs = {
'train': path.join( Params.CheckpointDir, Params.ModelName, 'training.log' ),
'valid': path.join( Params.CheckpointDir, Params.ModelName, 'validation.log' )
}
# Do not clear log files if we are continuing training a previous model
if not Params.isTrue('ContinueTrain'):
for log_filename in self.logs.values():
print( f"Clearing log file {log_filename}" )
if path.exists(log_filename):
os.remove(log_filename)
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 Inceptionv3_Net(n_classes, **kwargs):
if Params.isTrue('UseSqueezeExcitation'):
return Inception3(n_classes, inception_blocks=WrappedBlocks, **kwargs)
else:
return Inception3(n_classes, **kwargs)
return WrappedInceptionBlock
class Conv2d_NoBatchNorm(nn.Module):
def __init__(self, in_channels, out_channels, **kwargs):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(in_channels, out_channels, bias=False, **kwargs),
nn.ReLU(inplace=True))
def forward(self, x):
return self.net(x)
if Params.isTrue('BatchNorm'):
WrappedBlocks = [BasicConv2d]
else:
WrappedBlocks = [Conv2d_NoBatchNorm]
WrappedBlocks += [
se_wrap_block(BlockClass) for BlockClass in
[InceptionA, InceptionB, InceptionC, InceptionD, InceptionE]
]
WrappedBlocks += [InceptionAux]
def Inceptionv3_Net(n_classes, **kwargs):
if Params.isTrue('UseSqueezeExcitation'):
return Inception3(n_classes, inception_blocks=WrappedBlocks, **kwargs)
else:
return Inception3(n_classes, **kwargs)
