Update pages/15_CNN.py

pages/15_CNN.py

@@ -79,64 +79,67 @@ model = CNN().to(device)
 criterion = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
 
-# Lists to store losses
-train_losses = []
-test_losses = []
-
-# Train the model
-total_step = len(train_loader)
-for epoch in range(num_epochs):
-    train_loss = 0
-    for i, (images, labels) in enumerate(train_loader):
-        images = images.to(device)
-        labels = labels.to(device)
-
-        # Forward pass
-        outputs = model(images)
-        loss = criterion(outputs, labels)
-
-        # Backward and optimize
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        train_loss += loss.item()
-
-    train_loss /= total_step
-    train_losses.append(train_loss)
-    st.write(f'Epoch [{epoch+1}/{num_epochs}], Loss: {train_loss:.4f}')
-
-    # Test the model
-    model.eval()
-    with torch.no_grad():
-        test_loss = 0
-        correct = 0
-        total = 0
-        for images, labels in test_loader:
+# Button to start training
+if st.button("Start Training"):
+    # Lists to store losses
+    train_losses = []
+    test_losses = []
+
+    # Train the model
+    total_step = len(train_loader)
+    for epoch in range(num_epochs):
+        train_loss = 0
+        for i, (images, labels) in enumerate(train_loader):
             images = images.to(device)
             labels = labels.to(device)
+
+            # Forward pass
             outputs = model(images)
             loss = criterion(outputs, labels)
-            test_loss += loss.item()
-            _, predicted = torch.max(outputs.data, 1)
-            total += labels.size(0)
-            correct += (predicted == labels).sum().item()
-
-        test_loss /= len(test_loader)
-        test_losses.append(test_loss)
-        accuracy = 100 * correct / total
-        st.write(f'Test Loss: {test_loss:.4f}, Accuracy: {accuracy:.2f}%')
-    model.train()
-
-# Plotting the loss
-fig, ax = plt.subplots()
-ax.plot(range(1, num_epochs + 1), train_losses, label='Train Loss')
-ax.plot(range(1, num_epochs + 1), test_losses, label='Test Loss')
-ax.set_xlabel('Epoch')
-ax.set_ylabel('Loss')
-ax.set_title('Training and Test Loss')
-ax.legend()
-st.pyplot(fig)
-
-# Save the model checkpoint
-torch.save(model.state_dict(), 'cnn_model.pth')
+
+            # Backward and optimize
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+
+        train_loss /= total_step
+        train_losses.append(train_loss)
+        st.write(f'Epoch [{epoch+1}/{num_epochs}], Loss: {train_loss:.4f}')
+
+        # Test the model
+        model.eval()
+        with torch.no_grad():
+            test_loss = 0
+            correct = 0
+            total = 0
+            for images, labels in test_loader:
+                images = images.to(device)
+                labels = labels.to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                test_loss += loss.item()
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+            test_loss /= len(test_loader)
+            test_losses.append(test_loss)
+            accuracy = 100 * correct / total
+            st.write(f'Test Loss: {test_loss:.4f}, Accuracy: {accuracy:.2f}%')
+        model.train()
+
+    # Plotting the loss
+    fig, ax = plt.subplots()
+    ax.plot(range(1, num_epochs + 1), train_losses, label='Train Loss')
+    ax.plot(range(1, num_epochs + 1), test_losses, label='Test Loss')
+    ax.set_xlabel('Epoch')
+    ax.set_ylabel('Loss')
+    ax.set_title('Training and Test Loss')
+    ax.legend()
+    st.pyplot(fig)
+
+    # Save the model checkpoint
+    torch.save(model.state_dict(), 'cnn_model.pth')
+    st.write("Model training completed and saved as 'cnn_model.pth'")