Argument | Type | Value | Description |
---|---|---|---|
pretrained_model | – | bert-base-cased | A small BERT pre-trained model |
framework | - | PyTorch | Machine learning framework |
num_train_epochs | int | 5 | The number of epochs |
max_seq_length | int | 128 | Maximum sequence length the model support |
train_batch_size | int | 8 | Training batch size |
eval_batch_size | int | 8 | Evaluation batch size |
fp16 | bool | False | fp16 mode used |
dataloader_num_workers | int | 0 | Number of worker processed |
learning_rate | float | 4e−5 | Learning rate for training |
n_gpu | int | 1 | Number of GPU used |
optimizer | str | “AdamW” | Optimizer used |
save_eval_checkpoints | bool | False | Save model for each epcoh |