refine training

haoyu-he · Nov 6, 2023 · a591af9 · a591af9
1 parent 05e31b5
commit a591af9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/config.py b/config.py
@@ -18,7 +18,8 @@ class Config:
 
     batch = 32
     epoch = 5
-    lr = 3e-3
+    lr_lstm = 2e-3
+    lr_transformer = 2e-4
 
     train_size = 0.8
 

diff --git a/train.py b/train.py
@@ -55,7 +55,7 @@
 
 criterion = torch.nn.CrossEntropyLoss().to(config.device)
 parameters = list(encoder.parameters()) + list(emb_layer.parameters()) + list(decoder.parameters())
-optimizer = torch.optim.Adam(params=parameters, lr=config.lr)
+optimizer = torch.optim.Adam(params=parameters, lr=config.lr_lstm if  args.model == 'lstm' else config.lr_transformer)
 
 # training
 print('---Training---')