115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
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 sys
|
|
|
|
|
|
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])
|
|
print(epochs)
|
|
|
|
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)
|
|
|
|
# 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')
|
|
|
|
if __name__ == "__main__":
|
|
main()
|