Add advanced training metrics and confusion matrix logging/visualization

app.py

@@ -17,6 +17,9 @@ from dotenv import load_dotenv
 from chatbot_utils import ask_receipt_chatbot
 import time
 from tensorboard.backend.event_processing import event_accumulator
+from torch.utils.tensorboard import SummaryWriter
+import matplotlib.pyplot as plt
+from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -204,35 +207,73 @@ def main():
 
     if st.button("Start Training"):
         epochs = 10
+        num_classes = 3  # Example: 3 classes for confusion matrix
         losses = []
+        val_losses = []
         accuracies = []
         progress = st.progress(0)
-        chart = st.line_chart({"Loss": [], "Accuracy": []})
+        chart = st.line_chart({"Loss": [], "Val Loss": [], "Accuracy": []})
+
+        writer = SummaryWriter("logs")
 
         for epoch in range(epochs):
             # Simulate training
             loss = np.exp(-epoch/5) + np.random.rand() * 0.05
+            val_loss = loss + np.random.rand() * 0.02
             acc = 1 - loss + np.random.rand() * 0.02
             losses.append(loss)
+            val_losses.append(val_loss)
             accuracies.append(acc)
-            chart.add_rows({"Loss": [loss], "Accuracy": [acc]})
+            chart.add_rows({"Loss": [loss], "Val Loss": [val_loss], "Accuracy": [acc]})
             progress.progress((epoch+1)/epochs)
-            st.write(f"Epoch {epoch+1}: Loss={loss:.4f}, Accuracy={acc:.4f}")
+            st.write(f"Epoch {epoch+1}: Loss={loss:.4f}, Val Loss={val_loss:.4f}, Accuracy={acc:.4f}")
+
+            # Log to TensorBoard
+            writer.add_scalar("loss", loss, epoch)
+            writer.add_scalar("val_loss", val_loss, epoch)
+            writer.add_scalar("accuracy", acc, epoch)
+
+            # Simulate predictions and labels for confusion matrix
+            y_true = np.random.randint(0, num_classes, 100)
+            y_pred = y_true.copy()
+            # Add some noise to predictions
+            y_pred[np.random.choice(100, 10, replace=False)] = np.random.randint(0, num_classes, 10)
+            cm = confusion_matrix(y_true, y_pred, labels=range(num_classes))
+
+            # Plot and log confusion matrix as image
+            fig, ax = plt.subplots()
+            disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
+            disp.plot(ax=ax)
+            plt.close(fig)
+            writer.add_figure("confusion_matrix", fig, epoch)
+
             time.sleep(0.5)
 
+        writer.close()
         st.success("Training complete!")
-        st.line_chart({"Loss": losses, "Accuracy": accuracies})
+        st.line_chart({"Loss": losses, "Val Loss": val_losses, "Accuracy": accuracies})
 
-    logdir = "logs"  # or the actual path to your logs
+        # Show last confusion matrix in Streamlit
+        st.subheader("Confusion Matrix (Last Epoch)")
+        fig, ax = plt.subplots()
+        disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
+        disp.plot(ax=ax)
+        st.pyplot(fig)
 
+    logdir = "logs"
     if os.path.exists(logdir) and os.listdir(logdir):
         ea = event_accumulator.EventAccumulator(logdir)
         ea.Reload()
-        if 'loss' in ea.Tags()['scalars']:
-            losses = [s.value for s in ea.Scalars('loss')]
-            st.line_chart(losses)
-        else:
-            st.info("No 'loss' scalar found in TensorBoard logs.")
+        scalars = ea.Tags()['scalars']
+        for tag in ['loss', 'val_loss', 'accuracy']:
+            if tag in scalars:
+                values = [s.value for s in ea.Scalars(tag)]
+                st.line_chart({tag: values})
+        # Show confusion matrix images if available
+        if 'confusion_matrix' in ea.Tags()['images']:
+            st.subheader("TensorBoard Confusion Matrices")
+            for img in ea.Images('confusion_matrix'):
+                st.image(img.encoded_image_string)
     else:
         st.info("No TensorBoard logs found. Please upload logs to the 'logs' directory.")
 

requirements.txt

@@ -31,3 +31,5 @@ pydantic>=2.0.0
 openai
 streamlit
 plotly==5.18.0
+matplotlib
+scikit-learn