Daniel Voigt Godoy - Deep Learning with PyTorch Step-by-Step A Beginner’s Guide-leanpub

The code should look like this:
Mini-Batch
1 def _mini_batch(self, validation=False):
2 # The mini-batch can be used with both loaders
3 # The argument `validation` defines which loader and
4 # corresponding step function are going to be used
5 if validation:
6 data_loader = self.val_loader
7 step_fn = self.val_step_fn
8 else:
9 data_loader = self.train_loader
10 step_fn = self.train_step_fn
11
12 if data_loader is None:
13 return None
14
15 # Once the data loader and step function are set, this is the
16 # same mini-batch loop we had before
17 mini_batch_losses = []
18 for x_batch, y_batch in data_loader:
19 x_batch = x_batch.to(self.device)
20 y_batch = y_batch.to(self.device)
21
22 mini_batch_loss = step_fn(x_batch, y_batch)
23 mini_batch_losses.append(mini_batch_loss)
24
25 loss = np.mean(mini_batch_losses)
26
27 return loss
28
29 setattr(StepByStep, '_mini_batch', _mini_batch)
Moreover, if the user decides not to provide a validation loader, it will retain its
initial None value from the constructor method. If that’s the case, we don’t have a
corresponding loss to compute, and it returns None instead (line 13 in the snippet
above).
What’s left to do? The training loop, of course! This is similar to our Model Training
V5 in Chapter 2, but we can make it more flexible, taking the number of epochs and
188 | Chapter 2.1: Going Classy