Update pages/2_LinearRegression.py

pages/2_LinearRegression.py

@@ -3,47 +3,83 @@ import numpy as np
 import matplotlib.pyplot as plt
 import torch
 import torch.nn as nn
+import torch.optim as optim
 
-# Set a seed for reproducibility
-torch.manual_seed(59)
+# Streamlit app title
+st.title('Simple Linear Regression with PyTorch')
+
+# Sidebar sliders for noise and number of data points
+noise_level = st.sidebar.slider('Noise Level', 0.0, 1.0, 0.1, step=0.01)
+num_points = st.sidebar.slider('Number of Data Points', 10, 100, 50, step=5)
+num_epochs = st.sidebar.slider('Number of Epochs', 10, 500, 100, step=10)
+learning_rate = st.sidebar.slider('Learning Rate', 0.001, 0.1, 0.01, step=0.001)
 
-# Define the Linear Model
-class LinearModel(nn.Module):
-    def __init__(self, in_features, out_features):
-        super(LinearModel, self).__init__()
-        self.linear = nn.Linear(in_features, out_features)
+# Generate data
+np.random.seed(0)
+x = np.linspace(0, 10, num_points)
+y = 2 * x + 1 + noise_level * np.random.randn(num_points)
+
+# Convert data to PyTorch tensors
+x_tensor = torch.tensor(x, dtype=torch.float32).view(-1, 1)
+y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1)
+
+# Define the linear regression model
+class LinearRegressionModel(nn.Module):
+    def __init__(self):
+        super(LinearRegressionModel, self).__init__()
+        self.linear = nn.Linear(1, 1)
     
     def forward(self, x):
         return self.linear(x)
 
-# Instantiate the model
-model = LinearModel(1, 1)
+model = LinearRegressionModel()
 
-# Print model weight and bias
-print(f'Model weight: {model.linear.weight.item()}')
-print(f'Model bias: {model.linear.bias.item()}')
+# Define the loss function and the optimizer
+criterion = nn.MSELoss()
+optimizer = optim.SGD(model.parameters(), lr=learning_rate)
 
-# Streamlit app title
-st.title('Interactive Scatter Plot with Noise and Number of Data Points')
+# Train the model
+losses = []
+for epoch in range(num_epochs):
+    model.train()
+    optimizer.zero_grad()
+    outputs = model(x_tensor)
+    loss = criterion(outputs, y_tensor)
+    loss.backward()
+    optimizer.step()
+    losses.append(loss.item())
 
-# Sidebar sliders for noise and number of data points
-noise_level = st.sidebar.slider('Noise Level', 0.0, 1.0, 0.1, step=0.01)
-num_points = st.sidebar.slider('Number of Data Points', 10, 100, 50, step=5)
+# Get the final model parameters
+slope = model.linear.weight.item()
+intercept = model.linear.bias.item()
 
-# Generate data
-np.random.seed(59)
-x = np.linspace(0, 10, num_points).reshape(-1, 1).astype(np.float32)
-with torch.no_grad():
-    x_tensor = torch.tensor(x)
-    y_tensor = model(x_tensor)
-    y = y_tensor.numpy().flatten() + noise_level * np.random.randn(num_points)
-
-# Create scatter plot
+# Make predictions
+model.eval()
+y_pred_tensor = model(x_tensor)
+y_pred = y_pred_tensor.detach().numpy()
+
+# Create scatter plot with regression line
 fig, ax = plt.subplots()
-ax.scatter(x, y, alpha=0.6)
+ax.scatter(x, y, alpha=0.6, label='Data points')
+ax.plot(x, y_pred, color='red', label='Regression line')
 ax.set_title('Scatter Plot with Noise and Number of Data Points')
 ax.set_xlabel('X-axis')
 ax.set_ylabel('Y-axis')
+ax.legend()
+
+# Display slope and intercept in Streamlit
+st.write(f"**Slope:** {slope}")
+st.write(f"**Intercept:** {intercept}")
+
+# Display scatter plot in Streamlit
+st.pyplot(fig)
+
+# Plot training loss
+fig_loss, ax_loss = plt.subplots()
+ax_loss.plot(range(num_epochs), losses)
+ax_loss.set_title('Training Loss')
+ax_loss.set_xlabel('Epoch')
+ax_loss.set_ylabel('Loss')
 
-# Display plot in Streamlit
-st.pyplot(fig)
+# Display training loss plot in Streamlit
+st.pyplot(fig_loss)