Update pages/13_FFNN.py

pages/13_FFNN.py

@@ -37,6 +37,7 @@ def load_data():
 
 # Function to train the network
 def train_network(net, trainloader, criterion, optimizer, epochs):
+    loss_values = []
     for epoch in range(epochs):
         running_loss = 0.0
         for i, data in enumerate(trainloader, 0):
@@ -47,8 +48,11 @@ def train_network(net, trainloader, criterion, optimizer, epochs):
             loss.backward()
             optimizer.step()
             running_loss += loss.item()
-        st.write(f'Epoch {epoch + 1}: loss {running_loss / len(trainloader):.3f}')
+        epoch_loss = running_loss / len(trainloader)
+        loss_values.append(epoch_loss)
+        st.write(f'Epoch {epoch + 1}: loss {epoch_loss:.3f}')
     st.write('Finished Training')
+    return loss_values
 
 # Function to test the network
 def test_network(net, testloader):
@@ -61,7 +65,9 @@ def test_network(net, testloader):
             _, predicted = torch.max(outputs.data, 1)
             total += labels.size(0)
             correct += (predicted == labels).sum().item()
-    st.write(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')
+    accuracy = 100 * correct / total
+    st.write(f'Accuracy of the network on the 10000 test images: {accuracy:.2f}%')
+    return accuracy
 
 # Load the data
 trainloader, testloader = load_data()
@@ -85,11 +91,21 @@ st.write('\n' * 10)
 
 # Train the network
 if st.sidebar.button('Train Network'):
-    train_network(net, trainloader, criterion, optimizer, epochs)
+    loss_values = train_network(net, trainloader, criterion, optimizer, epochs)
+    
+    # Plot the loss values
+    plt.figure(figsize=(10, 5))
+    plt.plot(range(1, epochs + 1), loss_values, marker='o')
+    plt.title('Training Loss Over Epochs')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.grid(True)
+    st.pyplot(plt)
 
 # Test the network
 if st.sidebar.button('Test Network'):
-    test_network(net, testloader)
+    accuracy = test_network(net, testloader)
+    st.write(f'Test Accuracy: {accuracy:.2f}%')
 
 # Visualize some test results
 def imshow(img):