Refactor DNN model to support variable number of

app.py

@@ -9,30 +9,30 @@ from sklearn.model_selection import train_test_split
 import csv
 import io
 
-# from joblib import load, dump
-
 
 # Define the DNN model
 class DNN(nn.Module):
-    def __init__(self, input_size, hidden_size, output_size):
+    def __init__(self, input_size, hidden_size, output_size, num_hidden_layers):
         super(DNN, self).__init__()
         self.fc1 = nn.Linear(input_size, hidden_size)
         self.relu1 = nn.ReLU()
-        self.fc2 = nn.Linear(hidden_size, hidden_size)
-        self.relu2 = nn.ReLU()
+        self.hidden_layers = nn.ModuleList()
+        for _ in range(num_hidden_layers):
+            self.hidden_layers.append(nn.Linear(hidden_size, hidden_size))
+            self.hidden_layers.append(nn.ReLU())
         self.fc3 = nn.Linear(hidden_size, output_size)
 
     def forward(self, x):
         x = self.fc1(x)
         x = self.relu1(x)
-        x = self.fc2(x)
-        x = self.relu2(x)
+        for layer in self.hidden_layers:
+            x = layer(x)
         x = self.fc3(x)
         return x
 
 
 # Load the model
-model = DNN(input_size=6, hidden_size=256, output_size=1)
+model = DNN(input_size=6, hidden_size=64, output_size=1, num_hidden_layers=32)
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
@@ -71,23 +71,19 @@ def generate(
     prediction = model(input_data)
     return {"prediction": prediction.item()}
 
+    @app.post("/train")
+    async def train(
+        trainDatafile: UploadFile = File(...),
+        testDatafile: UploadFile = File(...),
+        epochs: int = 100,
+    ):
+        global model
 
-@app.post("/train")
-async def train(
-    trainDatafile: UploadFile = File(...),
-    testDatafile: UploadFile = File(...),
-    epochs: int = 100,
-):
-    global model
-
-    contents1 = await trainDatafile.read()
-    train_data = pd.read_csv(io.StringIO(contents1.decode("utf-8")))
-
-    contents2 = await testDatafile.read()
-    test_data = pd.read_csv(io.StringIO(contents2.decode("utf-8")))
+        contents1 = await trainDatafile.read()
+        train_data = pd.read_csv(io.StringIO(contents1.decode("utf-8")))
 
-    # Load the training and testing data
-    # test_data = pd.read_csv("dataset/agricultural_yield_test.csv")
+        contents2 = await testDatafile.read()
+        test_data = pd.read_csv(io.StringIO(contents2.decode("utf-8")))
 
     # Convert data to numpy arrays
     X_train = train_data.drop("Yield_kg_per_hectare", axis=1).values
@@ -128,7 +124,9 @@ async def train(
                 y_test.cpu().detach().numpy(), predictions.cpu().detach().numpy()
             )
         )
-        print(f"Epoch: {epoch+1}, RMSE: {float(rmse)}")
+        print(
+            f"Epoch: {epoch+1}, RMSE: {float(rmse)}, Loss: {float(np.sqrt(loss.cpu().detach().numpy()))}"
+        )
         rmseList.append(float(rmse))
 
     torch.save(model.state_dict(), "model_weights.pth")