74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
import torch
|
|
from torch import nn
|
|
from torch.optim import Adam
|
|
|
|
from ml_lib.modules.blocks import ConvModule
|
|
from ml_lib.modules.utils import LightningBaseModule, Flatten
|
|
|
|
|
|
class BinaryClassifier(LightningBaseModule):
|
|
def test_step(self, *args, **kwargs):
|
|
pass
|
|
|
|
def test_epoch_end(self, outputs):
|
|
pass
|
|
|
|
@classmethod
|
|
def name(cls):
|
|
return cls.__name__
|
|
|
|
def configure_optimizers(self):
|
|
return Adam(params=self.parameters(), lr=self.hparams.lr)
|
|
|
|
def training_step(self, batch_xy, batch_nb, *args, **kwargs):
|
|
batch_x, batch_y = batch_xy
|
|
y = self(batch_x)
|
|
loss = self.criterion(y, batch_y)
|
|
return dict(loss=loss)
|
|
|
|
def validation_step(self, batch_xy, **kwargs):
|
|
batch_x, batch_y = batch_xy
|
|
y = self(batch_y)
|
|
val_loss = self.criterion(y, batch_y)
|
|
return dict(val_loss=val_loss)
|
|
|
|
def validation_epoch_end(self, outputs):
|
|
overall_val_loss = torch.mean(torch.stack([output['val_loss'] for output in outputs]))
|
|
return dict(log=dict(
|
|
mean_val_loss=overall_val_loss)
|
|
)
|
|
|
|
def __init__(self, hparams):
|
|
super(BinaryClassifier, self).__init__(hparams)
|
|
|
|
self.criterion = nn.BCELoss()
|
|
|
|
# Additional parameters
|
|
self.in_shape = self.hparams.in_shape
|
|
|
|
# Model Modules
|
|
self.conv_1 = ConvModule(self.in_shape, 32, 3, conv_stride=2, **self.hparams.module_paramters)
|
|
self.conv_2 = ConvModule(self.conv_1.shape, 64, 5, conv_stride=2, **self.hparams.module_paramters)
|
|
self.conv_3 = ConvModule(self.conv_2.shape, 128, 7, conv_stride=2, **self.hparams.module_paramters)
|
|
|
|
self.flat = Flatten(self.conv_3.shape)
|
|
self.full_1 = nn.Linear(self.flat.shape, 32, self.hparams.bias)
|
|
self.full_2 = nn.Linear(self.full_1.out_features, self.full_1.out_features // 2, self.hparams.bias)
|
|
|
|
self.activation = self.hparams.activation()
|
|
self.full_out = nn.Linear(self.full_2.out_features, 1, self.hparams.bias)
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, batch, **kwargs):
|
|
tensor = self.conv_1(batch)
|
|
tensor = self.conv_2(tensor)
|
|
tensor = self.conv_3(tensor)
|
|
tensor = self.flat(tensor)
|
|
tensor = self.full_1(tensor)
|
|
tensor = self.activation(tensor)
|
|
tensor = self.full_2(tensor)
|
|
tensor = self.activation(tensor)
|
|
tensor = self.full_out(tensor)
|
|
tensor = self.sigmoid(tensor)
|
|
return tensor
|