Add DNN model and training code

app.py

@@ -1,6 +1,8 @@
-from fastapi import FastAPI, UploadFile, File
+import torch
+import torch.nn as nn
+import torch.optim as optim
 import numpy as np
-from sklearn.neural_network import MLPRegressor
+from fastapi import FastAPI, UploadFile, File
 from sklearn.model_selection import KFold
 from sklearn.metrics import mean_squared_error
 import csv
@@ -8,8 +10,24 @@ import io
 
 from joblib import load, dump
 
+
+# Define the DNN model
+class DNN(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size):
+        super(DNN, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_size, output_size)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return x
+
+
 # Load the model
-model = load("model.joblib")
+model = DNN(input_size=4, hidden_size=16, output_size=1)
 
 # Create a new FastAPI app instance
 app = FastAPI(docs_url="/", redoc_url="/new_redoc")
@@ -24,8 +42,11 @@ def generate(
     yearBuilt: float,
 ):
     global model
-    prediction = model.predict([[squareFeet, bedrooms, bathrooms, yearBuilt]])
-    return {"output": prediction[0]}
+    input_data = torch.tensor(
+        [[squareFeet, bedrooms, bathrooms, yearBuilt]], dtype=torch.float32
+    )
+    prediction = model(input_data)
+    return {"output": prediction.item()}
 
 
 @app.post("/train")
@@ -48,6 +69,14 @@ async def train(file: UploadFile = File(...)):
     y = data_np[:, -1]
     y = np.ravel(y)
 
+    # Convert data to torch tensors
+    X = torch.tensor(X, dtype=torch.float32)
+    y = torch.tensor(y, dtype=torch.float32)
+
+    # Define loss function and optimizer
+    criterion = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
     # Fit the model
     kf = KFold(n_splits=4)
     accuracies = []
@@ -56,10 +85,18 @@ async def train(file: UploadFile = File(...)):
         X_train, X_test = X[train_index], X[test_index]
         y_train, y_test = y[train_index], y[test_index]
 
-        model.fit(X_train, y_train)
+        optimizer.zero_grad()
+
+        # Forward pass
+        outputs = model(X_train)
+        loss = criterion(outputs, y_train.unsqueeze(1))
+
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
 
-        predictions = model.predict(X_test)
-        rmse = np.sqrt(mean_squared_error(y_test, predictions))
+        predictions = model(X_test)
+        rmse = np.sqrt(mean_squared_error(y_test, predictions.detach().numpy()))
         accuracies.append(rmse)
 
     average_rmse = sum(accuracies) / len(accuracies)

docker-compose.yml

@@ -6,6 +6,9 @@ services:
       - "7860:7860"
     deploy:
       resources:
+        limits:
+          memory: 512M
+          shm_size: 2G
         reservations:
           devices:
             - driver: nvidia

requirements.txt

@@ -28,3 +28,4 @@ uvicorn==0.17.6
 uvloop==0.19.0
 watchgod==0.8.2
 websockets==12.0
+torch