File size: 2,420 Bytes
84bfd88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
from typing import Tuple, Any, Union
import torch
from torch import nn
import numpy as np
class BaseModel(nn.Module):
def __init__(self):
super(BaseModel, self).__init__()
# set device (gpu 0 or 1 if available or cpu)
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# make empty param dict
self.param_dict = {}
# make optimizer options dict
self.optimizer_dict = {
"adam": torch.optim.Adam,
"nadam": torch.optim.NAdam,
"adamax": torch.optim.Adamax,
"adamw": torch.optim.AdamW,
"sgd": torch.optim.SGD,
}
# make loss options dict
self.loss_dict = {
"mse": nn.MSELoss,
"l1": nn.L1Loss,
"smoothl1": nn.SmoothL1Loss,
"huber": nn.HuberLoss,
"cel": nn.CrossEntropyLoss, # Suitable for classification tasks
"bcel": nn.BCELoss, # Suitable for classification tasks
}
# make activation function options dictionary
self.activation_dict = {
"relu": nn.ReLU,
"swish": nn.Hardswish,
"leaky_relu": nn.LeakyReLU,
"elu": nn.ELU,
"selu": nn.SELU,
}
# make tokenizer placeholder
self.tokenizer = None
# create history dictionary
self.history = {
"train_loss": [],
"on_diag_loss": [],
"off_diag_loss": [],
"validation_loss": [],
"learning_rate": [],
}
# create early stopping params
self.count = 0
def print_config(self) -> None:
print("[BT]: Current parameter config:")
print(self.param_dict)
def early_stopping(self, patience: int) -> bool:
# count every epoch that's worse than the best for patience times
if len(self.history["validation_loss"]) > patience:
best_loss = min(self.history["validation_loss"])
if self.history["validation_loss"][-1] > best_loss:
self.count += 1
else:
self.count = 0
if self.count >= patience:
if self.param_dict["verbose"] is True:
print("[BT]: Early stopping")
return True
return False
|