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)