Model Training

This commit is contained in:
Si11ium
2020-05-03 18:00:51 +02:00
parent 8a97f59906
commit e4f6506a4b
9 changed files with 167 additions and 105 deletions

View File

@@ -1,24 +1,23 @@
from argparse import Namespace
import torch
from torch import nn
from torch.optim import Adam
from torchvision.transforms import Compose, ToTensor
from ml_lib.audio_toolset.audio_io import Melspectogram, NormalizeLocal
from ml_lib.modules.blocks import ConvModule
from ml_lib.modules.utils import LightningBaseModule, Flatten
from ml_lib.modules.utils import LightningBaseModule, Flatten, BaseModuleMixin_Dataloaders
class BinaryClassifier(LightningBaseModule):
def test_step(self, *args, **kwargs):
pass
def test_epoch_end(self, outputs):
pass
class BinaryClassifier(BaseModuleMixin_Dataloaders, LightningBaseModule):
@classmethod
def name(cls):
return cls.__name__
def configure_optimizers(self):
return Adam(params=self.parameters(), lr=self.hparams.lr)
return Adam(params=self.parameters(), lr=self.params.lr)
def training_step(self, batch_xy, batch_nb, *args, **kwargs):
batch_x, batch_y = batch_xy
@@ -26,11 +25,11 @@ class BinaryClassifier(LightningBaseModule):
loss = self.criterion(y, batch_y)
return dict(loss=loss)
def validation_step(self, batch_xy, **kwargs):
def validation_step(self, batch_xy, batch_idx, *args, **kwargs):
batch_x, batch_y = batch_xy
y = self(batch_y)
y = self(batch_x)
val_loss = self.criterion(y, batch_y)
return dict(val_loss=val_loss)
return dict(val_loss=val_loss, batch_idx=batch_idx)
def validation_epoch_end(self, outputs):
overall_val_loss = torch.mean(torch.stack([output['val_loss'] for output in outputs]))
@@ -41,22 +40,36 @@ class BinaryClassifier(LightningBaseModule):
def __init__(self, hparams):
super(BinaryClassifier, self).__init__(hparams)
self.criterion = nn.BCELoss()
# Dataset and Dataloaders
# =============================================================================
# Transforms
transforms = Compose([Melspectogram(), ToTensor(), NormalizeLocal()])
# Datasets
from datasets.binar_masks import BinaryMasksDataset
self.dataset = Namespace(
**dict(
train_dataset=BinaryMasksDataset(self.params.root, setting='train', transforms=transforms),
val_dataset=BinaryMasksDataset(self.params.root, setting='devel', transforms=transforms),
test_dataset=BinaryMasksDataset(self.params.root, setting='test', transforms=transforms),
)
)
# Model Paramters
# =============================================================================
# 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.in_shape = self.dataset.train_dataset.sample_shape
self.criterion = nn.BCELoss()
# Modules
self.conv_1 = ConvModule(self.in_shape, 32, 3, conv_stride=2, **self.params.module_kwargs)
self.conv_2 = ConvModule(self.conv_1.shape, 64, 5, conv_stride=2, **self.params.module_kwargs)
self.conv_3 = ConvModule(self.conv_2.shape, 128, 7, conv_stride=2, **self.params.module_kwargs)
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.full_1 = nn.Linear(self.flat.shape, 32, self.params.bias)
self.full_2 = nn.Linear(self.full_1.out_features, self.full_1.out_features // 2, self.params.bias)
self.activation = self.hparams.activation()
self.full_out = nn.Linear(self.full_2.out_features, 1, self.hparams.bias)
self.activation = self.params.activation()
self.full_out = nn.Linear(self.full_2.out_features, 1, self.params.bias)
self.sigmoid = nn.Sigmoid()
def forward(self, batch, **kwargs):