|
- from .model import Model
- import torch
- from .batch import PredictionsBatch
- from typing import Callable
- from types import FunctionType
-
-
- class TrainLoop(object):
- def __init__(self, model: Model, lr: float = 0.001,
- loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = \
- torch.nn.functional.binary_cross_entropy_with_logits,
- batch_size: int = 100) -> None:
-
- if not isinstance(model, Model):
- raise TypeError('model must be an instance of Model')
-
- lr = float(lr)
-
- if not isinstance(loss, FunctionType):
- raise TypeError('loss must be a function')
-
- batch_size = int(batch_size)
-
- self.model = model
- self.lr = lr
- self.loss = loss
- self.batch_size = batch_size
-
- self.opt = None
-
- self.build()
-
- def build(self) -> None:
- opt = torch.optim.Adam(self.model.parameters(), lr=self.lr)
- self.opt = opt
-
- def run_epoch(self):
- pred = self.model(None)
- batch = PredictionsBatch(pred, batch_size=self.batch_size)
- n = len(list(iter(batch)))
- loss_sum = 0
- for i in range(n):
- self.opt.zero_grad()
- pred = self.model(None)
- batch = PredictionsBatch(pred, batch_size=self.batch_size, shuffle=True)
- seed = torch.rand(1).item()
- rng_state = torch.get_rng_state()
- torch.manual_seed(seed)
- it = iter(batch)
- torch.set_rng_state(rng_state)
- for k in range(i):
- _ = next(it)
- (input, target) = next(it)
- loss = self.loss(input, target)
- loss.backward()
- self.opt.step()
- loss_sum += loss.detach().cpu().item()
- return loss_sum
-
-
- def train(self, max_epochs):
- best_loss = None
- best_epoch = None
- for i in range(max_epochs):
- loss = self.run_epoch()
- if best_loss is None or loss < best_loss:
- best_loss = loss
- best_epoch = i
- return loss, best_loss, best_epoch
|
