import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset import pandas as pd from sklearn.model_selection import train_test_split from sklearn.preprocessing import LabelEncoder import torch.nn.functional as F import mlflow import mlflow.sklearn import sys mlflow.set_tracking_uri("http://localhost:5000") mlflow.set_experiment("s464914") device = ( "cuda" if torch.cuda.is_available() else "cpu" ) class Model(nn.Module): def __init__(self, input_features=54, hidden_layer1=25, hidden_layer2=30, output_features=8): super().__init__() self.fc1 = nn.Linear(input_features,output_features) self.bn1 = nn.BatchNorm1d(hidden_layer1) # Add batch normalization self.fc2 = nn.Linear(hidden_layer1, hidden_layer2) self.bn2 = nn.BatchNorm1d(hidden_layer2) # Add batch normalization self.out = nn.Linear(hidden_layer2, output_features) def forward(self, x): x = F.relu(self.fc1(x)) # Apply batch normalization after first linear layer #x = F.relu(self.bn2(self.fc2(x))) # Apply batch normalization after second linear layer #x = self.out(x) return x def main(): epochs = int(sys.argv[1]) forest_train = pd.read_csv('forest_train.csv') forest_val = pd.read_csv('forest_val.csv') print(forest_train.head()) X_train = forest_train.drop(columns=['Cover_Type']).values y_train = forest_train['Cover_Type'].values X_val = forest_val.drop(columns=['Cover_Type']).values y_val = forest_val['Cover_Type'].values # Initialize model, loss function, and optimizer model = Model().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # Convert to PyTorch tensors X_train = torch.tensor(X_train, dtype=torch.float32).to(device) y_train = torch.tensor(y_train, dtype=torch.long).to(device) X_val = torch.tensor(X_val, dtype=torch.float32).to(device) y_val = torch.tensor(y_val, dtype=torch.long).to(device) # Create DataLoader train_loader = DataLoader(list(zip(X_train, y_train)), batch_size=64, shuffle=True) val_loader = DataLoader(list(zip(X_val, y_val)), batch_size=64) with mlflow.start_run() as run: # Training loop for epoch in range(epochs): model.train() # Set model to training mode running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) # Calculate training loss epoch_loss = running_loss / len(train_loader.dataset) # Validation model.eval() # Set model to evaluation mode val_running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) val_loss = criterion(outputs, labels) val_running_loss += val_loss.item() * inputs.size(0) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() # Calculate validation loss and accuracy val_epoch_loss = val_running_loss / len(val_loader.dataset) val_accuracy = correct / total print(f"Epoch {epoch+1}/{epochs}, " f"Train Loss: {epoch_loss:.4f}, " f"Val Loss: {val_epoch_loss:.4f}, " f"Val Accuracy: {val_accuracy:.4f}") torch.save(model.state_dict(), 'model.pth') mlflow.log_param("epochs", epochs) if __name__ == "__main__": main()