File size: 590 Bytes
d18eb09 |
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 |
import torch
MODEL_CONFIG = {
'VOCAB_SIZE': 50000,
'D_MODEL': 1024,
'N_HEADS': 32,
'D_FF': 4096,
'N_LAYERS': 32,
'MAX_SEQ_LEN': 512,
'BATCH_SIZE': 32,
'LEARNING_RATE': 1e-4,
'NUM_EPOCHS': 20,
'DEVICE': 'cuda' if torch.cuda.is_available() else 'cpu',
'WARMUP_STEPS': 4000,
'ADAM_EPSILON': 1e-8,
'WEIGHT_DECAY': 0.01,
'GRADIENT_ACCUMULATION_STEPS': 2,
'MAX_GRAD_NORM': 1.0,
'DROPOUT': 0.1,
}
TRAINING_CONFIG = {
'CHECKPOINT_SAVE_STEPS': 5000,
'LOGGING_STEPS': 100,
'EVAL_STEPS': 1000,
'SAVE_TOTAL_LIMIT': 5
}
|