Mini Neural Network

train.py

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from io import StringIO
from sklearn.preprocessing import LabelEncoder

# Define a dataset class
class CustomDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

# Load and preprocess data
data = """
is_fluffy,can_eat_you,can_fly,has_talons,class
1,0,0,0,dog
0,0,1,1,eagle
1,1,0,0,lion
0,0,1,0,plane
1,1,1,1,griffin
"""
df = pd.read_csv(StringIO(data))
features = df.drop('class', axis=1).values
labels = LabelEncoder().fit_transform(df['class'])

# Creating dataloaders
dataset = CustomDataset(
    torch.tensor(features, dtype=torch.float32),
    torch.tensor(labels, dtype=torch.long)
)
data_loader = DataLoader(dataset, batch_size=1, shuffle=True)

# Neural Network
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(4, 2)
        self.fc2 = nn.Linear(2, 5)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        print("fc1 output: ", x)
        x = self.fc2(x)
        return x

# Initialize the network, loss function, and optimizer
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# Training the model
def train(model, data_loader, criterion, optimizer, epochs=100):
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for data, target in data_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f'Epoch {epoch+1}, Loss: {total_loss/len(data_loader)}')

# Run training
train(model, data_loader, criterion, optimizer)

# Run each data point through the network and print the output
for data, target in data_loader:
    print("data: ", data)
    print("target: ", target)
    print("output: ", model(data))
    print("-"*80)

# print model weights
print("fc1 weight: ", model.fc1.weight)
print("fc1 bias: ", model.fc1.bias)
print("fc2 weight: ", model.fc2.weight)
print("fc2 bias: ", model.fc2.bias)