train_body_ann.py

import time
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, random_split
import pytorch_lightning as pl


# Define the dataset class
class CustomDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.int64)

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

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


# Define the LightningModule
class CustomModel(pl.LightningModule):
    def __init__(self, num_classes):
        super(CustomModel, self).__init__()
        self.layers = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(17 * 2, 30),
            nn.ReLU(),
            nn.BatchNorm1d(30),
            nn.Dropout(0.2),
            nn.Linear(30, 15),
            nn.ReLU(),
            nn.Linear(15, num_classes)
        )
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.layers(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        self.log('val_loss', loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())


# Load your dataset
dataset = 'data_body_12classes_24112023.csv'
X_dataset = np.loadtxt(dataset, delimiter=',', dtype='float32', usecols=list(range(1, (17 * 2) + 1)))
y_dataset = np.loadtxt(dataset, delimiter=',', dtype='int32', usecols=0)
print('X shape: ', X_dataset.shape)
print('y shape: ', y_dataset.shape)
# print('X samples: ', X_dataset[1:3])
print('y samples: ', y_dataset[-15:])
time.sleep(5)
# Convert to PyTorch tensors
X_dataset = torch.tensor(X_dataset)
y_dataset = torch.tensor(y_dataset)

# Create dataset and dataloaders
dataset = CustomDataset(X_dataset, y_dataset)
train_size = int(0.75 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=128)

# Initialize PyTorch Lightning model
NUM_CLASSES = 12
lightning_model = CustomModel(NUM_CLASSES)

# Initialize a trainer with EarlyStopping and ModelCheckpoint callbacks
trainer = pl.Trainer(max_epochs=1000,
                     callbacks=[
                         pl.callbacks.EarlyStopping(monitor='val_loss', patience=20),
                         pl.callbacks.ModelCheckpoint(
                             filename='body12cls_24112023_{epoch}-{val_loss:.2f}',
                             save_top_k=1,
                             verbose=True,
                             monitor='val_loss',
                             mode='min'
                         )
                     ])

# Train the model
trainer.fit(lightning_model, train_loader, val_loader)