提高训练速度并使用GPU进行训练

config.py

@@ -34,7 +34,7 @@ class Config:
 
     # System
     device: str = "cpu"
-    num_workers: int = 0
+    num_workers: int = 4
     save_dir: str = "./checkpoints"
     log_interval: int = 100
 

train.py

@@ -217,13 +217,14 @@ def main():
         learning_rate=config.learning_rate,
         weight_decay=config.weight_decay,
         warmup_steps=config.warmup_steps,
-        use_cpu=True,
+        use_cpu=False,
         report_to="none",
         save_total_limit=2,
         load_best_model_at_end=True,
         metric_for_best_model="f1",
         greater_is_better=True,
         dataloader_num_workers=config.num_workers,
+        fp16=True
     )
 
     # Data collator