From cbdd4288431e44bb77990a1c2c3f36c37c330bd4 Mon Sep 17 00:00:00 2001 From: yayoimizuha Date: Sun, 15 Oct 2023 23:12:25 +0900 Subject: [PATCH] apply PyTorch AMP --- finetune/facenet_transfer_learning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/finetune/facenet_transfer_learning.py b/finetune/facenet_transfer_learning.py index f06bb0f..7a25861 100644 --- a/finetune/facenet_transfer_learning.py +++ b/finetune/facenet_transfer_learning.py @@ -126,7 +126,7 @@ optimizer = Adam(params=[ {'params': model_gpu[1].parameters(), 'lr': 1e-3}, ]) -scaler = GradScaler() +scaler = GradScaler(init_scale=4096) # model, optimizer = optimize(model=model, optimizer=optimizer) scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=10, gamma=0.9) @@ -168,7 +168,7 @@ for epoch in range(epochs): scaler.scale(loss).backward() scaler.unscale_(optimizer=optimizer) - utils.clip_grad_norm_(model.parameters(), max_norm=.5) + utils.clip_grad_norm_(model.parameters(), max_norm=.1) # loss.backward() scaler.step(optimizer=optimizer) scaler.update()