perf(finetune): 添加权重衰减，设置随步长衰减为7轮

zjZSTU · Mar 31, 2020 · 38fc365 · 38fc365
1 parent 37d5204
commit 38fc365
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/py/finetune.py b/py/finetune.py
@@ -127,7 +127,7 @@ def train_model(data_loaders, model, criterion, optimizer, lr_scheduler, num_epo
 
     criterion = nn.CrossEntropyLoss()
     # optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
-    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)
+    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4, weight_decay=1e-4)
     lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=15, gamma=0.1)
 
     best_model = train_model(data_loaders, model, criterion, optimizer, lr_scheduler, device=device, num_epochs=50)