import os
import time
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import seaborn as sns
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW, Adam
from torch.optim.lr_scheduler import (CosineAnnealingLR,
CosineAnnealingWarmRestarts,
StepLR,
ExponentialLR)
from .meter import Meter
from .dataset import ECGDataset
from .models import *
from .config import Config, seed_everything
class Trainer:
def __init__(self, net, lr, batch_size, num_epochs):
self.net = net.to(config.device)
self.num_epochs = num_epochs
self.criterion = nn.CrossEntropyLoss()
self.optimizer = AdamW(self.net.parameters(), lr=lr)
self.scheduler = CosineAnnealingLR(self.optimizer, T_max=num_epochs, eta_min=5e-6)
self.best_loss = float('inf')
self.phases = ['train', 'val']
self.dataloaders = {
phase: get_dataloader(phase, batch_size) for phase in self.phases
}
self.train_df_logs = pd.DataFrame()
self.val_df_logs = pd.DataFrame()
def _train_epoch(self, phase):
print(f"{phase} mode | time: {time.strftime('%H:%M:%S')}")
self.net.train() if phase == 'train' else self.net.eval()
meter = Meter()
meter.init_metrics()
for i, (data, target) in enumerate(self.dataloaders[phase]):
data = data.to(config.device)
target = target.to(config.device)
output = self.net(data)
loss = self.criterion(output, target)
if phase == 'train':
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
meter.update(output, target, loss.item())
metrics = meter.get_metrics()
metrics = {k:v / i for k, v in metrics.items()}
df_logs = pd.DataFrame([metrics])
confusion_matrix = meter.get_confusion_matrix()
if phase == 'train':
self.train_df_logs = pd.concat([self.train_df_logs, df_logs], axis=0)
else:
self.val_df_logs = pd.concat([self.val_df_logs, df_logs], axis=0)
# show logs
print('{}: {}, {}: {}, {}: {}, {}: {}, {}: {}'
.format(*(x for kv in metrics.items() for x in kv))
)
fig, ax = plt.subplots(figsize=(5, 5))
cm_ = ax.imshow(confusion_matrix, cmap='hot')
ax.set_title('Confusion matrix', fontsize=15)
ax.set_xlabel('Actual', fontsize=13)
ax.set_ylabel('Predicted', fontsize=13)
plt.colorbar(cm_)
plt.show()
return loss
def run(self):
for epoch in range(self.num_epochs):
self._train_epoch(phase='train')
with torch.no_grad():
val_loss = self._train_epoch(phase='val')
self.scheduler.step()
if val_loss < self.best_loss:
self.best_loss = val_loss
print('\nNew checkpoint\n')
self.best_loss = val_loss
torch.save(self.net.state_dict(), f"best_model_epoc{epoch}.pth")
if __name__ == '__main__':
# init config and set random seed
config = Config()
seed_everything(config.seed)
# init model
#model = RNNAttentionModel(1, 64, 'lstm', False)
#model = RNNModel(1, 64, 'lstm', True)
model = CNN(num_classes=5, hid_size=128)
# start train
trainer = Trainer(net=model, lr=1e-3, batch_size=96, num_epochs=30)
trainer.run()
# write logs
train_logs = trainer.train_df_logs
train_logs.columns = ["train_"+ colname for colname in train_logs.columns]
val_logs = trainer.val_df_logs
val_logs.columns = ["val_"+ colname for colname in val_logs.columns]
logs = pd.concat([train_logs,val_logs], axis=1)
logs.reset_index(drop=True, inplace=True)
logs = logs.loc[:, [
'train_loss', 'val_loss',
'train_accuracy', 'val_accuracy',
'train_f1', 'val_f1',
'train_precision', 'val_precision',
'train_recall', 'val_recall']
]
print(logs.head())
logs.to_csv('cnn.csv', index=False)
Datasets
Models
