|
|
|
|
|
|
|
|
|
import pdb |
|
from tqdm import tqdm |
|
from collections import defaultdict |
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
class Trainer(nn.Module): |
|
"""Helper class to train a deep network. |
|
Overload this class `forward_backward` for your actual needs. |
|
|
|
Usage: |
|
train = Trainer(net, loader, loss, optimizer) |
|
for epoch in range(n_epochs): |
|
train() |
|
""" |
|
|
|
def __init__(self, net, loader, loss, optimizer): |
|
nn.Module.__init__(self) |
|
self.net = net |
|
self.loader = loader |
|
self.loss_func = loss |
|
self.optimizer = optimizer |
|
|
|
def iscuda(self): |
|
return next(self.net.parameters()).device != torch.device("cpu") |
|
|
|
def todevice(self, x): |
|
if isinstance(x, dict): |
|
return {k: self.todevice(v) for k, v in x.items()} |
|
if isinstance(x, (tuple, list)): |
|
return [self.todevice(v) for v in x] |
|
|
|
if self.iscuda(): |
|
return x.contiguous().cuda(non_blocking=True) |
|
else: |
|
return x.cpu() |
|
|
|
def __call__(self): |
|
self.net.train() |
|
|
|
stats = defaultdict(list) |
|
|
|
for iter, inputs in enumerate(tqdm(self.loader)): |
|
inputs = self.todevice(inputs) |
|
|
|
|
|
self.optimizer.zero_grad() |
|
|
|
loss, details = self.forward_backward(inputs) |
|
if torch.isnan(loss): |
|
raise RuntimeError("Loss is NaN") |
|
|
|
self.optimizer.step() |
|
|
|
for key, val in details.items(): |
|
stats[key].append(val) |
|
|
|
print(" Summary of losses during this epoch:") |
|
mean = lambda lis: sum(lis) / len(lis) |
|
for loss_name, vals in stats.items(): |
|
N = 1 + len(vals) // 10 |
|
print(f" - {loss_name:20}:", end="") |
|
print( |
|
f" {mean(vals[:N]):.3f} --> {mean(vals[-N:]):.3f} (avg: {mean(vals):.3f})" |
|
) |
|
return mean(stats["loss"]) |
|
|
|
def forward_backward(self, inputs): |
|
raise NotImplementedError() |
|
|
