Implement MLP like in the Paper (?)

mlp_train.py

+import torch.nn as nn
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import trange, tqdm
+from sklearn.model_selection import train_test_split
+import time
+import matplotlib.pyplot as plt
+
+PICKLE_PATH = "dataset_pandas/temperature.pkl"
+
+
+##### HYPERPARAMETERS #####
+EPOCHS = 500
+BATCH_SIZE = 16
+CRITERION = nn.BCELoss()
+OPTIMIZER = torch.optim.Adam
+LEARNING_RATE = 0.01
+GROWTH_RATE = 16
+DROP_RATE = 0.5
+
+
+class MLP(nn.Module):
+    def __init__(self):
+        super(MLP, self).__init__()
+        self.block1 = nn.Sequential(
+            nn.Linear(20, GROWTH_RATE), nn.BatchNorm1d(GROWTH_RATE), nn.ReLU()
+        )
+        self.block2 = nn.Sequential(
+            nn.Linear(GROWTH_RATE, GROWTH_RATE), nn.BatchNorm1d(GROWTH_RATE), nn.ReLU()
+        )
+        self.block3 = nn.Sequential(
+            nn.Linear(2 * GROWTH_RATE, GROWTH_RATE),
+            nn.BatchNorm1d(GROWTH_RATE),
+            nn.ReLU(),
+        )
+        self.block4 = nn.Sequential(
+            nn.Linear(3 * GROWTH_RATE, GROWTH_RATE),
+            nn.BatchNorm1d(GROWTH_RATE),
+            nn.ReLU(),
+        )
+        self.block5 = nn.Sequential(
+            nn.Linear(4 * GROWTH_RATE, GROWTH_RATE),
+            nn.BatchNorm1d(GROWTH_RATE),
+            nn.ReLU(),
+        )
+        self.block6 = nn.Sequential(
+            nn.Linear(5 * GROWTH_RATE, GROWTH_RATE),
+            nn.BatchNorm1d(GROWTH_RATE),
+            nn.ReLU(),
+        )
+        self.output = nn.Linear(GROWTH_RATE, 1)
+
+    def forward(self, x):
+        x1 = self.block1(x)
+        x2 = self.block2(x1)
+        x3 = self.block3(torch.cat([x1, x2], dim=1))
+        x4 = self.block4(torch.cat([x1, x2, x3], dim=1))
+        x5 = self.block5(torch.cat([x1, x2, x3, x4], dim=1))
+        x6 = self.block6(torch.cat([x1, x2, x3, x4, x5], dim=1))
+        y = self.output(x6)
+        y = torch.sigmoid(y)
+        return y
+
+
+def prepare_data() -> tuple[np.ndarray, np.ndarray]:
+    # Load data
+    df = pd.read_pickle(PICKLE_PATH)
+
+    X = df.drop(columns=["alarm"]).to_numpy()
+    y = df["alarm"].to_numpy()
+
+    assert X.shape[1] == 20, "Number of features should be 20"
+    assert y.shape[0] == X.shape[0], "Number of labels should match number of samples"
+
+    return X, y
+
+
+def train_model(X: np.ndarray, y: np.ndarray):
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Splitting the data into training and testing sets
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
+
+    # Setting up the data loader
+    train_loader = torch.utils.data.DataLoader(
+        list(zip(X_train, y_train)), batch_size=BATCH_SIZE, shuffle=True
+    )
+
+    # Define model
+    model = MLP().to(device)
+
+    model.dropout = nn.Dropout(DROP_RATE)
+
+    # Define loss function
+    criterion = CRITERION
+
+    # Define optimizer
+    optimizer = OPTIMIZER(model.parameters(), lr=LEARNING_RATE)
+
+    # Define a Scheduler
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="min", factor=0.1, patience=10, verbose=True, eps=1e-8
+    )
+
+    # Train model
+    train_losses = []
+    test_losses = []
+    for epoch in range(EPOCHS):
+        with tqdm(train_loader, unit="batch") as t:
+            for data, target in t:
+                t.set_description(f"Epoch {str(epoch).rjust(5)}")
+
+                # Move data to device
+                data, target = data.to(device), target.to(device)
+
+                # Zero the gradients
+                optimizer.zero_grad()
+
+                output = model(data.float())
+
+                # Calculate loss
+                loss = criterion(output, target.float().view(-1, 1))
+
+                # Backpropagation
+                loss.backward()
+
+                # Update weights
+                optimizer.step()
+
+                # Display loss
+                t.set_postfix(train_loss=f"{loss.item():.4f}")
+
+            scheduler.step(loss)
+
+            train_losses.append(loss.item())
+            # Evaluate model on test set
+            y_pred = (
+                model(torch.tensor(X_test).float().to(device)).cpu().detach().numpy()
+            )
+            test_loss = criterion(
+                torch.tensor(y_pred).float(), torch.tensor(y_test).float().view(-1, 1)
+            )
+            print(f"Test loss: {test_loss.item():.4f}")
+            test_losses.append(test_loss.item())
+
+    # Plot losses
+    plt.plot(train_losses, label="Train loss")
+    plt.plot(test_losses, label="Test loss")
+    plt.legend()
+    plt.show()
+
+    return model
+
+
+def main():
+    X, y = prepare_data()
+    model = train_model(X, y)
+
+
+if __name__ == "__main__":
+    main()