105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
from abc import ABC
|
|
|
|
import torch
|
|
import pandas as pd
|
|
|
|
from ml_lib.modules.util import LightningBaseModule
|
|
from util.loss_mixin import LossMixin
|
|
from util.optimizer_mixin import OptimizerMixin
|
|
|
|
|
|
class TrainMixin:
|
|
|
|
def training_step(self, batch_xy, batch_nb, *args, **kwargs):
|
|
assert isinstance(self, LightningBaseModule)
|
|
batch_files, batch_x, batch_y = batch_xy
|
|
y = self(batch_x).main_out
|
|
|
|
if self.params.loss == 'focal_loss_rob':
|
|
labels_one_hot = torch.nn.functional.one_hot(batch_y, num_classes=5)
|
|
loss = self.__getattribute__(self.params.loss)(y, labels_one_hot)
|
|
else:
|
|
loss = self.__getattribute__(self.params.loss)(y, batch_y.long())
|
|
|
|
return dict(loss=loss)
|
|
|
|
def training_epoch_end(self, outputs):
|
|
assert isinstance(self, LightningBaseModule)
|
|
keys = list(outputs[0].keys())
|
|
|
|
summary_dict = {f'mean_{key}': torch.mean(torch.stack([output[key]
|
|
for output in outputs]))
|
|
for key in keys if 'loss' in key}
|
|
summary_dict.update(epoch=self.current_epoch)
|
|
self.log_dict(summary_dict)
|
|
|
|
|
|
class ValMixin:
|
|
|
|
def validation_step(self, batch_xy, batch_idx, *args, **kwargs):
|
|
assert isinstance(self, LightningBaseModule)
|
|
batch_files, batch_x, batch_y = batch_xy
|
|
model_out = self(batch_x)
|
|
y = model_out.main_out
|
|
|
|
val_loss = self.ce_loss(y, batch_y.long())
|
|
|
|
self.metrics.update(y, batch_y) # torch.argmax(y, -1), batch_y)
|
|
|
|
return dict(val_loss=val_loss, batch_files=batch_files,
|
|
batch_idx=batch_idx, y=y, batch_y=batch_y)
|
|
|
|
def validation_epoch_end(self, outputs, *_, **__):
|
|
assert isinstance(self, LightningBaseModule)
|
|
summary_dict = dict()
|
|
|
|
keys = list(outputs[0].keys())
|
|
summary_dict.update({f'mean_{key}': torch.mean(torch.stack([output[key]
|
|
for output in outputs]))
|
|
for key in keys if 'loss' in key}
|
|
)
|
|
# Sklearn Scores
|
|
additional_scores = self.additional_scores(outputs)
|
|
summary_dict.update(**additional_scores)
|
|
|
|
pl_metrics, pl_images = self.metrics.compute_and_prepare()
|
|
self.metrics.reset()
|
|
summary_dict.update(**pl_metrics)
|
|
summary_dict.update(epoch=self.current_epoch)
|
|
|
|
self.log_dict(summary_dict, on_epoch=True)
|
|
|
|
for name, image in pl_images.items():
|
|
self.logger.log_image(name, image, step=self.global_step)
|
|
pass
|
|
|
|
|
|
class TestMixin:
|
|
|
|
def test_step(self, batch_xy, batch_idx, *_, **__):
|
|
assert isinstance(self, LightningBaseModule)
|
|
batch_files, batch_x, batch_y = batch_xy
|
|
model_out = self(batch_x)
|
|
y = model_out.main_out
|
|
return dict(batch_files=batch_files, batch_idx=batch_idx, y=y)
|
|
|
|
def test_epoch_end(self, outputs, *_, **__):
|
|
assert isinstance(self, LightningBaseModule)
|
|
|
|
y_arg_max = torch.argmax(outputs[0]['y'])
|
|
|
|
pd.DataFrame(data=dict(filenames=outputs[0]['batch_files'], predtiction=y_arg_max))
|
|
|
|
# No logging, just inference.
|
|
# self.log_dict(summary_dict, on_epoch=True)
|
|
|
|
|
|
class CombinedModelMixins(LossMixin,
|
|
TrainMixin,
|
|
ValMixin,
|
|
TestMixin,
|
|
OptimizerMixin,
|
|
LightningBaseModule,
|
|
ABC):
|
|
pass
|