faknow.train

faknow.train.base_gnn_trainer

class faknow.train.base_gnn_trainer.BaseGNNTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]

Bases: BaseTrainer

Base trainer for GNN models, which inherits from BaseTrainer and modifies the evaluate method.

evaluate(loader: DataLoader)[source]

Evaluate model performance on testing or validation data.

Parameters:

loader (DataLoader) – pyg data to evaluate, where each batch data is torch_geometric.data.Batch and each sample data in a batch is torch_geometric.data.Data

Returns:

evaluation metrics

Return type:

Dict[str, float]

faknow.train.cafe_trainer

class faknow.train.cafe_trainer.CafeTrainer(model: CAFE, evaluator: Evaluator, detection_optimizer: Optimizer, similarity_optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]

Bases: BaseTrainer

Trainer for CAFE model with tow model, which inherits from BaseTrainer and modifies the ‘_train_epoch’ method.

__init__(model: CAFE, evaluator: Evaluator, detection_optimizer: Optimizer, similarity_optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]
Parameters:
  • model (CAFE) – the first faknow abstract model to train

  • evaluator (Evaluator) – faknow evaluator for evaluation

  • detection_optimizer (Optimizer) – pytorch optimizer for training of the detection model

  • similarity_optimizer (Optimizer) – pytorch optimizer for training of the similarity model

  • scheduler (_LRScheduler) – learning rate scheduler. Defaults=None.

  • clip_grad_norm (Dict[str, Any]) – key args for torch.nn.utils.clip_grad_norm_. Defaults=None.

  • device (str) – device to use. Defaults=’cuda:0’.

  • early_stopping (EarlyStopping) – early stopping for training. If None, no early stopping will be performed. Defaults=None.

faknow.train.dense_gnn_trainer

class faknow.train.dense_gnn_trainer.DenseGNNTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]

Bases: BaseGNNTrainer

Base trainer for GNN models with dense batch data, which inherits from BaseGNNTrainer and modifies the evaluate method.

evaluate(loader: DenseDataLoader)[source]

Evaluate model performance on testing or validation data.

Parameters:

loader (DenseDataLoader) – pyg dense data to evaluate, where each batch data is torch_geometric.data.Batch with all attributes stacked in a new dimension.

Returns:

evaluation metrics

Return type:

Dict[str, float]

faknow.train.pgd_trainer

class faknow.train.pgd_trainer.MFANTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]

Bases: BaseTrainer

Trainer for MFAN model with PGD defence, which inherits from BaseTrainer and modifies the ‘_train_epoch’ method.

faknow.train.trainer

class faknow.train.trainer.AbstractTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu')[source]

Bases: object

Abstract trainer for FaKnow, including training, validation, evaluation and save model. All trainers should inherit from this class.

model

faknow abstract model to train

Type:

AbstractModel

evaluator

faknow evaluator for evaluation

Type:

Evaluator

optimizer

pytorch optimizer for training

Type:

Optimizer

scheduler

learning rate scheduler.

Type:

_LRScheduler

clip_grad_norm

key args for torch.nn.utils.clip_grad_norm_.

Type:

Dict[str, Any]

device

device to use.

Type:

torch.device

__init__(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu')[source]
Parameters:
  • model (AbstractModel) – faknow abstract model to train

  • evaluator (Evaluator) – faknow evaluator for evaluation

  • optimizer (Optimizer) – pytorch optimizer for training

  • scheduler (_LRScheduler) – learning rate scheduler. Defaults=None.

  • clip_grad_norm (Dict[str, Any]) – key args for torch.nn.utils.clip_grad_norm_. Defaults=None.

  • device (str) – device to use. Defaults=’cpu’.

cpu(**kwargs)[source]

move model and data to cpu

Parameters:

**kwargs – other args for nn.Module.to

cuda(device='cuda:0', **kwargs)[source]

move model and data to cuda

Parameters:
  • device (str) – cuda device to use. Defaults=’cuda:0’.

  • **kwargs – other args for nn.Module.to

evaluate(loader: DataLoader)[source]

evaluate after training

Parameters:

loader (DataLoader) – data to evaluate

fit(train_loader: DataLoader, num_epochs: int, validate_loader: DataLoader | None = None, save=False, save_path=None)[source]

train model

Parameters:
  • train_loader (DataLoader) – training data

  • num_epochs (int) – number of epochs to train

  • validate_loader (DataLoader) – validation data. If None, no validation. Defaults=None.

  • save (bool) – whether to save model. Defaults=False.

  • save_path (str) – path to save model, if save=True. Defaults=None.

to(device: str, **kwargs)[source]

move model and data to device

Parameters:
  • device (str) – device to use

  • **kwargs – other args for nn.Module.to

class faknow.train.trainer.BaseTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]

Bases: AbstractTrainer

Base trainer for FaKnow, which inherits from AbstractTrainer and can be applied to most common tasks.

model

faknow abstract model to train

Type:

AbstractModel

evaluator

faknow evaluator for evaluation

Type:

Evaluator

optimizer

pytorch optimizer for training

Type:

Optimizer

scheduler

learning rate scheduler.

Type:

_LRScheduler

clip_grad_norm

key args for torch.nn.utils.clip_grad_norm_.

Type:

Dict[str, Any]

device

device to use.

Type:

torch.device

early_stopping

early stopping for training.

Type:

EarlyStopping

best_score

best validation score for saving best model.

Type:

float

best_epoch

best epoch for saving best model.

Type:

int

writer

tensorboard writer.

Type:

SummaryWriter

logger

logger for logging in console and file.

Type:

logging.Logger

__init__(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, clip_grad_norm: Dict[str, Any] | None = None, device='cpu', early_stopping: EarlyStopping | None = None)[source]
Parameters:
  • model (AbstractModel) – faknow abstract model to train

  • evaluator (Evaluator) – faknow evaluator for evaluation

  • optimizer (Optimizer) – pytorch optimizer for training

  • scheduler (_LRScheduler) – learning rate scheduler. Defaults=None.

  • clip_grad_norm (Dict[str, Any]) – key args for torch.nn.utils.clip_grad_norm_. Defaults=None.

  • device (str) – device to use. Defaults=’cpu’.

  • early_stopping (EarlyStopping) – early stopping for training. If None, no early stopping will be performed. Defaults=None.

evaluate(loader: DataLoader)[source]

Evaluate model performance on testing or validation data.

Parameters:

loader (DataLoader) – data to evaluate, where each batch data is a dict with key ‘label’

Returns:

evaluation metrics

Return type:

Dict[str, float]

fit(train_loader: DataLoader, num_epochs: int, validate_loader: DataLoader | None = None, save_best: bool | None = None, save_path: str | None = None)[source]

Fit the model, including training, validation and save model. Results will be logged in console, log file and tensorboard. Early stopping has higher priority than save_best if it is not None.

Parameters:
  • train_loader (DataLoader) – training data

  • num_epochs (int) – number of epochs to train

  • validate_loader (DataLoader) – validation data. If None, no validation. Defaults=None.

  • save_best (bool) – whether to save model with best validation score. If False, save the last epoch model. If None, do not save any model. Defaults=None.

  • save_path (str) – path to save model, if save_best is not None. If None, save in ‘./save/model_name-current_time.pth’. Defaults=None.

save(save_path: str | None = None)[source]

save the model

Parameters:

save_path (str) – path to save model. If None, save in ‘./save/model_name-current_time.pth’. Defaults=None.

faknow.train.trainer_gpu

class faknow.train.trainer_gpu.AbstractTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, loss_func: Callable | None = None)[source]

Bases: object

__init__(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, loss_func: Callable | None = None)[source]
evaluate(data: DataLoader)[source]

evaluate after training

fit(train_data: DataLoader, num_epochs: int, validate_data=None, save=False, save_path=None)[source]
class faknow.train.trainer_gpu.BaseTrainer(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, loss_func: Callable | None = None)[source]

Bases: AbstractTrainer

__init__(model: AbstractModel, evaluator: Evaluator, optimizer: Optimizer, scheduler: _LRScheduler | None = None, loss_func: Callable | None = None)[source]
evaluate(loader: DataLoader)[source]

evaluate after training

fit(train_loader: DataLoader, num_epochs: int, validate_loader: DataLoader | None = None, save=False, save_path: str | None = None)[source]
save(save_path: str)[source]

save the model