title: PyTorch: is there a definitive training loop similar to Keras' fit()?
body: <p>I'm coming over from Keras to PyTorch, and one of the surprising things I've found is that I'm supposed to implement my own training loop.</p>
  
  <p>In Keras, there is a de facto <a href="https://keras.io/models/model/#fit" rel="noreferrer"><code>fit()</code></a> function that: (1) runs gradient descent and (2) collects a history of metrics for loss and accuracy over both the training set and validation set.</p>
  
  <p>In PyTorch, it appears that the programmer needs to implement the training loop. Since I'm new to PyTorch, I don't know if my training loop implementation is correct. <strong>I just want to compare apples-to-apples loss and accuracy metrics with what I'm seeing in Keras.</strong></p>
  
  <p>I've already read through:</p>
  
  <ol>
  <li><p>the official <a href="https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html" rel="noreferrer">PyTorch 60-minute blitz</a>, where they provide <a href="https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#train-the-network" rel="noreferrer">a sample training loop</a>.</p></li>
  <li><p>official <a href="https://github.com/pytorch/tutorials" rel="noreferrer">PyTorch example code</a>, where I've found <a href="https://github.com/pytorch/tutorials/blob/master/beginner_source/examples_nn/two_layer_net_module.py" rel="noreferrer">the training loop placed in-line</a> with other code.</p></li>
  <li><p>the O'Reilly book <a href="http://shop.oreilly.com/product/0636920216032.do" rel="noreferrer">Programming PyTorch for Deep Learning</a> with <a href="https://github.com/falloutdurham/beginners-pytorch-deep-learning/blob/master/chapter3/Chapter%203.ipynb" rel="noreferrer">its own training loop</a>.</p></li>
  <li><p>Stanford CS230 <a href="https://github.com/cs230-stanford/cs230-code-examples/tree/master/pytorch" rel="noreferrer">sample code</a>.</p></li>
  <li><p>various blog posts (e.g. <a href="https://www.analyticsvidhya.com/blog/2018/02/pytorch-tutorial/" rel="noreferrer">here</a> and <a href="https://towardsdatascience.com/understanding-pytorch-with-an-example-a-step-by-step-tutorial-81fc5f8c4e8e" rel="noreferrer">here</a>).</p></li>
  </ol>
  
  <p>So I'm wondering: <strong>is there a definitive, universal training loop implementation that does the same thing and reports the same numbers as the Keras <code>fit()</code> function?</strong></p>
  
  <p>My points of frustration:</p>
  
  <ol>
  <li><p>Pulling data out of the dataloader is not consistent between image data and NLP data.</p></li>
  <li><p>Correctly computing loss and accuracy is not consistent in any sample code I've seen.</p></li>
  <li><p>Some code examples use <code>Variable</code>, while others do not.</p></li>
  <li><p>Unnecessarily detailed: moving data to/from the GPU; knowing when to call <code>zero_grad()</code>.</p></li>
  </ol>
  
  <p>For what it's worth, here is my current implementation. Are there any obvious bugs?</p>
  
  <pre class="lang-py prettyprint-override"><code>import time
  
  def train(model, optimizer, loss_fn, train_dl, val_dl, epochs=20, device='cuda'):
      '''
      Runs training loop for classification problems. Returns Keras-style
      per-epoch history of loss and accuracy over training and validation data.
  
      Parameters
      ----------
      model : nn.Module
          Neural network model
      optimizer : torch.optim.Optimizer
          Search space optimizer (e.g. Adam)
      loss_fn :
          Loss function (e.g. nn.CrossEntropyLoss())
      train_dl : 
          Iterable dataloader for training data.
      val_dl :
          Iterable dataloader for validation data.
      epochs : int
          Number of epochs to run
      device : string
          Specifies 'cuda' or 'cpu'
  
      Returns
      -------
      Dictionary
          Similar to Keras' fit(), the output dictionary contains per-epoch
          history of training loss, training accuracy, validation loss, and
          validation accuracy.
      '''
  
      print('train() called: model=%s, opt=%s(lr=%f), epochs=%d, device=%s\n' % \
            (type(model).__name__, type(optimizer).__name__,
             optimizer.param_groups[0]['lr'], epochs, device))
  
      history = {} # Collects per-epoch loss and acc like Keras' fit().
      history['loss'] = []
      history['val_loss'] = []
      history['acc'] = []
      history['val_acc'] = []
  
      start_time_sec = time.time()
  
      for epoch in range(epochs):
  
          # --- TRAIN AND EVALUATE ON TRAINING SET -----------------------------
          model.train()
          train_loss         = 0.0
          num_train_correct  = 0
          num_train_examples = 0
  
          for batch in train_dl:
  
              optimizer.zero_grad()
  
              x    = batch[0].to(device)
              y    = batch[1].to(device)
              yhat = model(x)
              loss = loss_fn(yhat, y)
  
              loss.backward()
              optimizer.step()
  
              train_loss         += loss.data.item() * x.size(0)
              num_train_correct  += (torch.max(yhat, 1)[1] == y).sum().item()
              num_train_examples += x.shape[0]
  
          train_acc   = num_train_correct / num_train_examples
          train_loss  = train_loss / len(train_dl.dataset)
  
  
          # --- EVALUATE ON VALIDATION SET -------------------------------------
          model.eval()
          val_loss       = 0.0
          num_val_correct  = 0
          num_val_examples = 0
  
          for batch in val_dl:
  
              x    = batch[0].to(device)
              y    = batch[1].to(device)
              yhat = model(x)
              loss = loss_fn(yhat, y)
  
              val_loss         += loss.data.item() * x.size(0)
              num_val_correct  += (torch.max(yhat, 1)[1] == y).sum().item()
              num_val_examples += y.shape[0]
  
          val_acc  = num_val_correct / num_val_examples
          val_loss = val_loss / len(val_dl.dataset)
  
  
          print('Epoch %3d/%3d, train loss: %5.2f, train acc: %5.2f, val loss: %5.2f, val acc: %5.2f' % \
                (epoch+1, epochs, train_loss, train_acc, val_loss, val_acc))
  
          history['loss'].append(train_loss)
          history['val_loss'].append(val_loss)
          history['acc'].append(train_acc)
          history['val_acc'].append(val_acc)
  
      # END OF TRAINING LOOP
  
  
      end_time_sec       = time.time()
      total_time_sec     = end_time_sec - start_time_sec
      time_per_epoch_sec = total_time_sec / epochs
      print()
      print('Time total:     %5.2f sec' % (total_time_sec))
      print('Time per epoch: %5.2f sec' % (time_per_epoch_sec))
  
      return history
  </code></pre>