Update pages/19_ResNet.py

pages/19_ResNet.py

@@ -66,14 +66,15 @@ def imshow(inp, title=None):
     std = np.array([0.229, 0.224, 0.225])
     inp = std * inp + mean
     inp = np.clip(inp, 0, 1)
-    plt.imshow(inp)
+    fig, ax = plt.subplots()
+    ax.imshow(inp)
     if title is not None:
-        plt.title(title)
-    plt.pause(0.001)
+        ax.set_title(title)
+    st.pyplot(fig)
 
 inputs, classes = next(iter(dataloaders['train']))
 out = torchvision.utils.make_grid(inputs)
-st.pyplot(imshow(out, title=[class_names[x] for x in classes]))
+imshow(out, title=[class_names[x] for x in classes])
 
 # Model Preparation Section
 st.markdown("""
@@ -169,20 +170,18 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
     
     # Plot training history
     epochs_range = range(num_epochs)
-    plt.figure(figsize=(10, 5))
-    plt.subplot(1, 2, 1)
-    plt.plot(epochs_range, train_loss_history, label='Training Loss')
-    plt.plot(epochs_range, val_loss_history, label='Validation Loss')
-    plt.legend(loc='upper right')
-    plt.title('Training and Validation Loss')
-
-    plt.subplot(1, 2, 2)
-    plt.plot(epochs_range, train_acc_history, label='Training Accuracy')
-    plt.plot(epochs_range, val_acc_history, label='Validation Accuracy')
-    plt.legend(loc='lower right')
-    plt.title('Training and Validation Accuracy')
-    plt.show()
-    st.pyplot(plt)
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
+    ax1.plot(epochs_range, train_loss_history, label='Training Loss')
+    ax1.plot(epochs_range, val_loss_history, label='Validation Loss')
+    ax1.legend(loc='upper right')
+    ax1.set_title('Training and Validation Loss')
+
+    ax2.plot(epochs_range, train_acc_history, label='Training Accuracy')
+    ax2.plot(epochs_range, val_acc_history, label='Validation Accuracy')
+    ax2.legend(loc='lower right')
+    ax2.set_title('Training and Validation Accuracy')
+
+    st.pyplot(fig)
     
     return model