import time
import numpy as np
import torch
from torch import nn, optim
from torch.utils import data
from torchvision import models
[docs]
class CNN:
"""CNN Trainer class.
This class is responsible for training a CNN model.
Parameters:
train_data (torchvision.datasets.ImageFolder): Training data.
validation_data (torchvision.datasets.ImageFolder): Validation data.
test_data (torchvision.datasets.ImageFolder): Test data.
batch_size (int): Batch size.
"""
def __init__(self, train_data, validation_data, test_data, batch_size):
"""The CNN Trainer class constructor."""
# Train data loader
self.train_loader = data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
# Validation data loader
self.validation_loader = data.DataLoader(validation_data, batch_size=batch_size, shuffle=False)
# Test data loader
self.test_loader = data.DataLoader(test_data, batch_size=batch_size, shuffle=False)
# Trainer device type
self.device = torch.device("cpu")
[docs]
def create_and_train_cnn(self, model_name, num_epochs, learning_rate, weight_decay, replications):
"""Create and train a CNN model.
Parameters:
model_name (str): Model name to be trained.
num_epochs (int): Number of epochs to be trained.
learning_rate (float): Learning rate to be used at train.
weight_decay (float): Weight decay to be used at train.
replications (int): Number of replications used at each trained model.
Returns:
(dict): A dict mapping keys to the:
* 'result_name': (str) Result name.
* 'acc_avg': (float) Average accuracy.
* 'iter_acc_max': (int) Iteration of maximum accuracy.
* 'duration': (float) Duration of training.
"""
begin = time.time()
sum = 0
acc_max = 0
for i in range(0, replications):
model = self.create_model(model_name)
optimizerSGD = self.create_optimizer(model, learning_rate, weight_decay)
criterionCEL = self.create_criterion()
self.train_model(model, self.train_loader, optimizerSGD, criterionCEL, model_name, num_epochs, learning_rate, weight_decay, i)
acc = self.evaluate_model(model, self.validation_loader)
sum = sum + acc
if acc > acc_max:
acc_max = acc
iter_acc_max = i
end = time.time()
acc_avg = sum / replications
duration = end - begin
result_name = f"{model_name}-{num_epochs}-{learning_rate}-{weight_decay}"
return result_name, acc_avg, iter_acc_max, duration
[docs]
def create_model(self, model_name):
"""Create a function to a CNN model to be trained.
Note:
At moment, the models available are: [VGG11, Alexnet, MobilenetV3Large].
Parameters:
model_name (str): CNN model name.
Returns:
(function): Function to CNN model selected.
"""
if (model_name=='VGG11'):
model = models.vgg11(weights='DEFAULT')
for param in model.parameters():
param.requires_grad = False
model.classifier[6] = nn.Linear(model.classifier[6].in_features,2)
return model
elif (model_name=='Alexnet'):
model = models.alexnet(weights='DEFAULT')
for param in model.parameters():
param.requires_grad = False
model.classifier[6] = nn.Linear(model.classifier[6].in_features,2)
return model
else: # 'if (model_name=='MobilenetV3Large' ou qualquer outra coisa para não dar erro)
model = models.mobilenet_v3_large(weights='DEFAULT')
for param in model.parameters():
param.requires_grad = False
model.classifier[3] = nn.Linear(model.classifier[3].in_features,2)
return model
[docs]
def create_optimizer(self, model, learning_rate, weight_decay):
"""Create an optimizer.
Parameters:
model (function): CNN function.
learning_rate (float): Learning rate
weight_decay (float): Weight decay
Returns:
(object): Optimizer object.
"""
update = []
for name,param in model.named_parameters():
if param.requires_grad == True:
update.append(param)
optimizerSGD = optim.SGD(update, lr=learning_rate, weight_decay=weight_decay)
return optimizerSGD
[docs]
def create_criterion(self):
"""Create a loss criterion.
Parameters:
None
Returns:
(object): Cross entropy loss object.
"""
criterionCEL = nn.CrossEntropyLoss()
return criterionCEL
[docs]
def train_model(self, model, train_loader, optimizer, criterion, model_name, num_epochs, learning_rate, weight_decay, replication):
"""Train a CNN model.
Train a CNN model and save it (PTH file) at 'models' directory.
Parameters:
model (function): Model function.
train_loader (DataLoader): Training data loader
optimizer (object): Optimizer object.
criterion (object): CEL object.
model_name (str): Model name.
num_epochs (int): Number of epochs.
learning_rate (float): Learning rate.
weight_decay (float): Weight decay.
replication (int): Replication.
Returns:
None
"""
model.to(self.device)
min_loss = 100
e_measures = []
for i in (range(1,num_epochs+1)):
train_loss = self.train_epoch(model, train_loader, optimizer, criterion)
if (train_loss < min_loss):
min_loss = train_loss
nome_arquivo = f"./models/{model_name}_{num_epochs}_{learning_rate}_{weight_decay}_{replication}.pth"
torch.save(model.state_dict(), nome_arquivo)
[docs]
def train_epoch(self, model, trainLoader, optimizer, criterion):
"""Train an epoch.
Parameters:
model (function): Model function.
trainLoader (DataLoader): Training data loader.
optimizer (object): Optimizer object.
criterion (object): CEL object.
Returns:
(float): Mean of losses.
"""
model.train()
losses = []
for X, y in trainLoader:
X = X.to(self.device)
y = y.to(self.device)
optimizer.zero_grad()
y_pred = model(X)
loss = criterion(y_pred, y)
loss.backward()
optimizer.step()
losses.append(loss.item())
model.eval()
return np.mean(losses)
## @fn evaluate_model
# @brief Evaluate a model
# @param model Model
# @param loader Data loader
# @return Accuracy
[docs]
def evaluate_model(self, model, loader):
"""Evaluate a model.
Parameters:
model (function): Model function.
loader (DataLoader): Data loader
Returns:
(float): Model (trained) accuracy.
"""
total = 0
correct = 0
for X, y in loader:
X, y = X.to(self.device), y.to(self.device)
output = model(X)
_, y_pred = torch.max(output, 1)
total += len(y)
correct += (y_pred == y).sum().cpu().data.numpy()
acc = correct/total
return acc