Spaces:
Sleeping
Sleeping
File size: 3,351 Bytes
b20c769 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import pandas as pd
import torch
from lightning import LightningModule
class Module(LightningModule):
def __init__(
self, network, loss, train_metrics, val_metrics, test_metrics, scheduler, optimizer
):
super().__init__()
self.model = network.instance
self.loss = loss
self.train_metrics = train_metrics
self.val_metrics = val_metrics
self.test_metrics = test_metrics
self.optimizer = optimizer
self.scheduler = scheduler
def forward(self, x):
return self.model(x)
def training_step(self, batch, batch_idx):
pred = self.model(batch)
loss = self.loss(pred, batch, average=True)
if "logits" in loss.keys():
loss.pop("logits")
for metric_name, metric_value in loss.items():
self.log(
f"train/{metric_name}",
metric_value,
sync_dist=True,
on_step=True,
on_epoch=True,
)
return loss
@torch.no_grad()
def validation_step(self, batch, batch_idx):
pred = self.model(batch)
loss = self.loss(pred, batch, average=True)
if "logits" in loss.keys():
self.val_metrics.update(loss["logits"])
loss.pop("logits")
else:
self.val_metrics.update(pred, batch)
for metric_name, metric_value in loss.items():
self.log(
f"val/{metric_name}",
metric_value,
sync_dist=True,
on_step=False,
on_epoch=True,
)
def on_validation_epoch_end(self):
metrics = self.val_metrics.compute()
for metric_name, metric_value in metrics.items():
self.log(
f"val/{metric_name}",
metric_value,
sync_dist=True,
on_step=False,
on_epoch=True,
)
@torch.no_grad()
def test_step(self, batch, batch_idx):
pred = self.model(batch)
loss = self.loss(pred, batch, average=True)
if "logits" in loss.keys():
self.test_metrics.update(loss["logits"])
loss.pop("logits")
else:
self.test_metrics.update(pred, batch)
def on_test_epoch_end(self):
metrics = self.test_metrics.compute()
if "results" in metrics.keys():
pd.DataFrame(metrics["results"]).T.to_csv("results.csv")
print("saving results dict")
metrics.pop("results")
for metric_name, metric_value in metrics.items():
self.log(
f"test/{metric_name}",
metric_value,
sync_dist=True,
on_step=False,
on_epoch=True,
)
def configure_optimizers(self):
optimizer = self.optimizer(params=self.parameters())
if self.scheduler is not None:
scheduler = self.scheduler(optimizer=optimizer)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val/loss",
"interval": "epoch",
"frequency": 1,
},
}
return {"optimizer": optimizer}
|