interactive plt in TL tutorial

Author: chsasank

beginner_source/transfer_learning_tutorial.py

@@ -45,6 +45,8 @@
 import copy
 import os
 
+plt.ion()   # interactive mode
+
 ######################################################################
 # Load Data
 # ---------
@@ -101,13 +103,16 @@
 # Let's visualize a few training images so as to understand the data
 # augmentations.
 
-def imshow(inp):
+def imshow(inp, title=None):
     """Imshow for Tensor."""
     inp = inp.numpy().transpose((1, 2, 0))
     mean = np.array([0.485, 0.456, 0.406])
     std = np.array([0.229, 0.224, 0.225])
     inp = std * inp + mean
     plt.imshow(inp)
+    if title is not None:
+        plt.title(title)
+    plt.pause(0.001)  # pause a bit so that plots are updated
 
 
 # Get a batch of training data
@@ -116,9 +121,7 @@ def imshow(inp):
 # Make a grid from batch
 out = torchvision.utils.make_grid(inputs)
 
-imshow(out)
-plt.title([dset_classes[x] for x in classes])
-plt.show()
+imshow(out, title=[dset_classes[x] for x in classes])
 
 
 ######################################################################
@@ -222,14 +225,12 @@ def visualize_model(model, num_images=5):
         else:
             inputs, labels = Variable(inputs), Variable(labels)
 
-        
         outputs = model(inputs)
         _, preds = torch.max(outputs.data, 1)
-        
+
         plt.figure()
-        imshow(inputs.cpu().data[0])
-        plt.title('pred: {}'.format(dset_classes[labels.data[0]]))
-        plt.show()
+        imshow(inputs.cpu().data[0],
+               title='pred: {}'.format(dset_classes[labels.data[0]]))
 
         if i == num_images - 1:
             break
@@ -337,3 +338,6 @@ def optim_scheduler_conv(model, epoch, init_lr=0.001, lr_decay_epoch=7):
 #
 
 visualize_model(model)
+
+plt.ioff()
+plt.show()