torch_em.classification.classification_trainer

 1import warnings
 2
 3import numpy as np
 4import torch
 5import torch_em
 6
 7
 8class ClassificationTrainer(torch_em.trainer.DefaultTrainer):
 9    def _validate_impl(self, forward_context):
10        self.model.eval()
11
12        loss_val = 0.0
13
14        # we use the syntax from sklearn.metrics to compute metrics
15        # over all the preditions
16        y_true, y_pred = [], []
17
18        with torch.no_grad():
19            for x, y in self.val_loader:
20                x, y = x.to(self.device), y.to(self.device)
21                with forward_context():
22                    pred, loss = self._forward_and_loss(x, y)
23                loss_val += loss.item()
24                y_true.append(y.detach().cpu().numpy())
25                y_pred.append(pred.max(1)[1].detach().cpu().numpy())
26
27        if torch.isnan(pred).any():
28            warnings.warn("Predictions are NaN")
29        loss_val /= len(self.val_loader)
30
31        y_true, y_pred = np.concatenate(y_true), np.concatenate(y_pred)
32        metric_val = self.metric(y_true, y_pred)
33
34        if self.logger is not None:
35            self.logger.log_validation(self._iteration, metric_val, loss_val, x, y, pred, y_true, y_pred)
36        return metric_val