tomokazu
/
helloproject-ai
mirror of https://github.com/yayoimizuha/helloproject-ai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

apply PyTorch AMP
continuous-integration/drone/push Build was killed Details

This commit is contained in:
yayoimizuha 2023-10-15 22:58:30 +09:00
parent 1d943d029b
commit af154fed8d
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
finetune/facenet_transfer_learning.py
View File

@ -1,6 +1,6 @@
from os import makedirs, environ
from torchinfo import summary
from torch.nn import Linear, Sequential, Dropout, CrossEntropyLoss, Identity, ReLU
from torch.nn import Linear, Sequential, Dropout, CrossEntropyLoss, Identity, ReLU, utils
from torchvision.transforms import Compose, RandomResizedCrop, RandomRotation, ToTensor, \
RandomHorizontalFlip, \
Resize, RandomAutocontrast, InterpolationMode, RandomErasing, \
@ -167,6 +167,8 @@ for epoch in range(epochs):
scaler.scale(loss).backward()
scaler.unscale_(optimizer=optimizer)
utils.clip_grad_norm_(model.parameters(), max_norm=.5)
# loss.backward()
scaler.step(optimizer=optimizer)
scaler.update()