EQLossTrainer¶
- class torcheeg.trainers.EQLossTrainer(model: Module, num_classes: int, class_frequency: List[int] | DataLoader, gamma: float = 0.9, lambd: float = 0.005, lr: float = 0.001, weight_decay: float = 0.0, devices: int = 1, accelerator: str = 'cpu', metrics: List[str] = ['accuracy'])[source][source]¶
A trainer class for EEG classification with Equalization (EQ) loss for imbalanced datasets.
Paper: Tan J, Wang C, Li B, et al. Equalization loss for long-tailedobject recognition[C]//Proceedings of the IEEE/CVF conference on computervision and pattern recognition. 2020: 11662-11671.
Related Project: https://github.com/tztztztztz/eql.detectron2
from torcheeg.models import CCNN from torcheeg.trainers import EQLossTrainer model = CCNN(in_channels=5, num_classes=2) trainer = EQLossTrainer(model, num_classes=2, class_frequency=[10, 20], gamma=0.9, lambd=0.005)
- Parameters:
model (nn.Module) – The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
num_classes (int) – The number of classes in the dataset.
class_frequency (List[int] or Dataloader) – The frequency of each class in the dataset. It can be a list of integers or a dataloader to calculate the frequency of each class in the dataset, traversing the data batch (
torch.utils.data.dataloader.DataLoader
,torch_geometric.loader.DataLoader
, etc). (default:None
)gamma (float) – The gamma parameter. (default:
0.9
)lambd (float) – The lambd parameter. (default:
0.005
)lr (float) – The learning rate. (default:
0.001
)weight_decay (float) – The weight decay. (default:
0.0
)devices (int) – The number of devices to use. (default:
1
)accelerator (str) – The accelerator to use. Available options are: ‘cpu’, ‘gpu’. (default:
"cpu"
)metrics (list of str) – The metrics to use. Available options are: ‘precision’, ‘recall’, ‘f1score’, ‘accuracy’, ‘matthews’, ‘auroc’, and ‘kappa’. (default:
["accuracy"]
)
- fit(train_loader: DataLoader, val_loader: DataLoader, max_epochs: int = 300, *args, **kwargs) Any [source]¶
- Parameters:
train_loader (DataLoader) – Iterable DataLoader for traversing the training data batch (
torch.utils.data.dataloader.DataLoader
,torch_geometric.loader.DataLoader
, etc).val_loader (DataLoader) – Iterable DataLoader for traversing the validation data batch (
torch.utils.data.dataloader.DataLoader
,torch_geometric.loader.DataLoader
, etc).max_epochs (int) – Maximum number of epochs to train the model. (default:
300
)