|
| 1 | +import pandas as pd |
| 2 | +import torch |
| 3 | +import torch.nn as nn |
| 4 | +import torch.nn.functional as F |
| 5 | +import torch.optim as optim |
| 6 | +from torch.utils.data import DataLoader, TensorDataset |
| 7 | + |
| 8 | + |
| 9 | +# Define the neural network model with Batch Normalization |
| 10 | +class NeuralNetwork(nn.Module): |
| 11 | + def __init__(self, input_channels, num_classes): |
| 12 | + super(NeuralNetwork, self).__init__() |
| 13 | + self.conv1 = nn.Conv2d(in_channels=input_channels, out_channels=30, kernel_size=(3, 3), stride=2) |
| 14 | + self.dropout1 = nn.Dropout(0.5) |
| 15 | + self.conv2 = nn.Conv2d(in_channels=30, out_channels=30, kernel_size=(3, 3), stride=2) |
| 16 | + self.dropout2 = nn.Dropout(0.5) |
| 17 | + self.flatten = nn.Flatten() |
| 18 | + self.fc1 = nn.Linear(30 * 6 * 6, 128) # Adjust based on your input size |
| 19 | + self.fc2 = nn.Linear(128, num_classes) |
| 20 | + |
| 21 | + def forward(self, x): |
| 22 | + x = F.relu(self.conv1(x)) |
| 23 | + x = self.dropout1(x) |
| 24 | + x = F.relu(self.conv2(x)) |
| 25 | + x = self.dropout2(x) |
| 26 | + x = self.flatten(x) |
| 27 | + x = F.relu(self.fc1(x)) |
| 28 | + x = F.softmax(self.fc2(x), dim=1) |
| 29 | + return x |
| 30 | + |
| 31 | + |
| 32 | +def fit(X_train: pd.DataFrame, y_train: pd.DataFrame, X_valid: pd.DataFrame, y_valid: pd.DataFrame): |
| 33 | + # Convert data to PyTorch tensors and reshape it for convolutional layers |
| 34 | + X_train_tensor = ( |
| 35 | + torch.tensor(X_train.values, dtype=torch.float32).view(-1, 1, 28, 28).to(device) |
| 36 | + ) # Reshape and move to GPU |
| 37 | + y_train_tensor = torch.tensor(y_train.values, dtype=torch.long).to(device) |
| 38 | + X_valid_tensor = torch.tensor(X_valid.values, dtype=torch.float32).view(-1, 1, 28, 28).to(device) |
| 39 | + y_valid_tensor = torch.tensor(y_valid.values, dtype=torch.long).to(device) |
| 40 | + |
| 41 | + # Create datasets and dataloaders |
| 42 | + train_dataset = TensorDataset(X_train_tensor, y_train_tensor) |
| 43 | + valid_dataset = TensorDataset(X_valid_tensor, y_valid_tensor) |
| 44 | + train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True) |
| 45 | + valid_loader = DataLoader(valid_dataset, batch_size=128, shuffle=False) |
| 46 | + |
| 47 | + # Initialize the model, loss function and optimizer |
| 48 | + model = NeuralNetwork(input_channels=1, num_classes=len(set(y_train))).to(device) |
| 49 | + criterion = nn.CrossEntropyLoss().to(device) |
| 50 | + optimizer = optim.Adam(model.parameters(), lr=0.0005) |
| 51 | + |
| 52 | + # Train the model |
| 53 | + num_epochs = 400 |
| 54 | + for epoch in range(num_epochs): |
| 55 | + model.train() |
| 56 | + for X_batch, y_batch in train_loader: |
| 57 | + optimizer.zero_grad() |
| 58 | + outputs = model(X_batch) |
| 59 | + loss = criterion(outputs, y_batch) |
| 60 | + loss.backward() |
| 61 | + optimizer.step() |
| 62 | + |
| 63 | + # Validate the model |
| 64 | + model.eval() |
| 65 | + valid_loss = 0 |
| 66 | + correct = 0 |
| 67 | + with torch.no_grad(): |
| 68 | + for X_batch, y_batch in valid_loader: |
| 69 | + outputs = model(X_batch) |
| 70 | + valid_loss += criterion(outputs, y_batch).item() |
| 71 | + _, predicted = torch.max(outputs, 1) |
| 72 | + correct += (predicted == y_batch).sum().item() |
| 73 | + |
| 74 | + accuracy = correct / len(valid_loader.dataset) |
| 75 | + print(f"Epoch {epoch+1}/{num_epochs}, Validation Accuracy: {accuracy:.4f}") |
| 76 | + |
| 77 | + return model |
| 78 | + |
| 79 | + |
| 80 | +def predict(model, X): |
| 81 | + X_tensor = torch.tensor(X.values, dtype=torch.float32).view(-1, 1, 28, 28).to(device) |
| 82 | + model.eval() |
| 83 | + with torch.no_grad(): |
| 84 | + outputs = model(X_tensor) |
| 85 | + _, predicted = torch.max(outputs, 1) |
| 86 | + return predicted.cpu().numpy().reshape(-1, 1) |
0 commit comments