Commit 522d9f0b authored by Oliver Wirth's avatar Oliver Wirth
Browse files

Add LSTM training routine

parent d7008417
%% Cell type:code id: tags:
``` python
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
from sklearn.preprocessing import StandardScaler
import torch
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
np.random.seed(42)
pd.set_option("display.max_columns", None)
```
%% Cell type:code id: tags:
``` python
# Data paths
data_path = Path('data')
train_path = data_path / 'train'
test_path = data_path / 'test'
model_path = Path('checkpoints')
# Load labels
train_labels = pd.read_csv(f'{train_path}_label.csv')
test_labels = pd.read_csv(f'{test_path}_label.csv')
# Merge train and test labels
all_labels = train_labels.append(test_labels, ignore_index = True)
all_labels = all_labels.dropna()
all_labels.ret = all_labels.ret.astype(int)
```
%% Cell type:code id: tags:
``` python
# Custom dataset per region
class RegionDataset(Dataset) :
def __init__(self, region_path, labels, transform = None) :
super().__init__()
self.region = region_path.name
self.transform = transform
self.dfs = []
for csv_path in region_path.iterdir() :
df = pd.read_csv(csv_path)
df = df.dropna()
df = df[(df.T != 0).any()]
label = labels[labels.file_name == csv_path.name].ret.values
if df.shape[0] > 0 and len(label) == 1 :
self.dfs.append((df, label[0]))
def __len__(self) :
return len(self.dfs)
def __getitem__(self, idx) :
df, label = self.dfs[idx]
if self.transform :
df = self.transform(df)
return df, label
```
%% Cell type:code id: tags:
``` python
# Preprocessing transformations
# Ideas: PCA, truncated SVD, MinMaxScaling
transform = transforms.Compose([
StandardScaler().fit_transform,
torch.FloatTensor
])
# Load train and test set
trainset = RegionDataset(train_path / '004', all_labels, transform = transform)
testset = RegionDataset(test_path / 'dummy', all_labels, transform = transform)
```
%% Cell type:code id: tags:
``` python
# Split into train and validation set
holdout = .2
n_val = int(len(trainset) * .2)
n_train = len(trainset) - n_val
trainset, valset = torch.utils.data.dataset.random_split(trainset, [n_train, n_val])
# Create data loader
batch_size = 1
trainloader = DataLoader(trainset, batch_size = batch_size, shuffle = True)
testloader = DataLoader(testset , batch_size = batch_size, shuffle = False)
valloader = DataLoader(valset , batch_size = batch_size, shuffle = False)
```
%% Cell type:code id: tags:
``` python
# NN hyperparameters
input_size = 75
lstm_size = 64
hidden_size = 32
output_size = 2
dropout = .5
# NN definition
class RNN(nn.Module) :
def __init__(self) :
super().__init__()
self.rec = nn.LSTM(input_size = input_size, hidden_size = lstm_size, batch_first = True)
self.clf = nn.Sequential(
nn.Dropout(dropout),
nn.Linear(lstm_size, hidden_size),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_size, output_size)
)
def forward(self, x) :
x, _ = self.rec(x)
x = x[:,-1,:]
x = self.clf(x)
return x
```
%% Cell type:code id: tags:
``` python
# Learning rate and starting epoch
lr = 2e-4
epoch = 0
# Init model and optimizer
model = RNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr = lr)
# CUDA
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(DEVICE)
```
%%%% Output: execute_result
RNN(
(rec): LSTM(75, 64, batch_first=True)
(clf): Sequential(
(0): Linear(in_features=64, out_features=32, bias=True)
(1): ReLU()
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=32, out_features=2, bias=True)
)
)
%% Cell type:code id: tags:
``` python
# Forward pass for a single batch
def forward_pass(model, samples, labels, criterion, optimizer = None) :
out = model(samples)
pred = out.argmax(dim = 1)
loss = criterion(out, labels)
if optimizer :
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = loss.item() * samples.size(0)
correct = (pred == labels).sum().item()
return loss, correct
# Forward pass for whole dataset
def train_loop(model, loader, criterion, optimizer = None) :
model.train(optimizer is not None)
running_loss = 0
running_correct = 0
for samples, labels in loader :
samples, labels = samples.to(DEVICE), labels.to(DEVICE)
loss, correct = forward_pass(model, samples, labels, criterion, optimizer)
running_loss += loss
running_correct += correct
return running_loss, running_correct
```
%% Cell type:code id: tags:
``` python
epochs = 20
loss_stats = pd.DataFrame(columns = ['train', 'val'])
acc_stats = pd.DataFrame(columns = ['train', 'val'])
# Train model
for epoch in range(epoch, epoch + epochs) :
print(f'=== Epoch {(epoch + 1):2} ===')
# Train loop
loss, correct = train_loop(model, trainloader, criterion, optimizer)
train_loss = loss / len(trainset)
train_acc = correct / len(trainset)
print(f' training loss = {train_loss:.4f}, acc = {train_acc:.4f}')
# Validation loop
loss, correct = train_loop(model, valloader, criterion)
val_loss = loss / len(valset)
val_acc = correct / len(valset)
print(f'validation loss = {val_loss:.4f}, acc = {val_acc:.4f}')
# Statistics
loss_stats = loss_stats.append({
'train': train_loss,
'val': val_loss
}, ignore_index = True)
acc_stats = acc_stats.append({
'train': train_acc,
'val': val_acc
}, ignore_index = True)
# Save best model
if loss_stats['val'].idxmin() == len(loss_stats) - 1 :
torch.save(model, model_path / f'model_{model.__class__.__name__}_e{epoch + 1}.pt')
epoch += 1
```
%% Cell type:code id: tags:
``` python
# Evaluate trained model on test set
running_loss, running_correct = train_loop(model, testloader, criterion)
n = len(testset)
print(f'Evaluation after {epoch} epochs: loss = {(running_loss / n):.4f}, acc = {(running_correct / n):.4f}')
```
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment