mirrors/PaddleClas
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/PaddleClas.git synced 2025-06-03 21:55:06 +08:00
Code Issues Projects Releases Wiki Activity

Accelerate dynamic graph amp training

Browse Source
This commit is contained in:
zhangbo9674 2021-12-20 06:36:56 +00:00
parent 7732a69f1b
commit b54ee04491
3 changed files with 13 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ppcls/engine/engine.py
View File

@ -250,6 +250,8 @@ class Engine(object):
self.scaler = paddle.amp.GradScaler(
init_loss_scaling=self.scale_loss,
use_dynamic_loss_scaling=self.use_dynamic_loss_scaling)
if self.config['AMP']['use_pure_fp16'] is True:
self.model = paddle.amp.decorate(models=self.model, level='O2')
self.max_iter = len(self.train_dataloader) - 1 if platform.system(
) == "Windows" else len(self.train_dataloader)

13
ppcls/engine/train/train.py
View File

@ -41,14 +41,15 @@ def train_epoch(engine, epoch_id, print_batch_step):
# image input
if engine.amp:
with paddle.amp.auto_cast(custom_black_list={
"flatten_contiguous_range", "greater_than"
}):
amp_level = 'O1'
if engine.config['AMP']['use_pure_fp16'] is True:
amp_level = 'O2'
with paddle.amp.auto_cast(custom_black_list={"flatten_contiguous_range", "greater_than"}, level=amp_level):
out = forward(engine, batch)
loss_dict = engine.train_loss_func(out, batch[1])
else:
out = forward(engine, batch)
loss_dict = engine.train_loss_func(out, batch[1])
loss_dict = engine.train_loss_func(out, batch[1])
# step opt and lr
if engine.amp:
@ -58,7 +59,7 @@ def train_epoch(engine, epoch_id, print_batch_step):
else:
loss_dict["loss"].backward()
engine.optimizer.step()
engine.optimizer.clear_grad()
engine.optimizer.clear_grad(set_to_zero=True)
engine.lr_sch.step()
# below code just for logging

5
ppcls/optimizer/optimizer.py
View File

@ -36,13 +36,15 @@ class Momentum(object):
momentum,
weight_decay=None,
grad_clip=None,
multi_precision=False):
multi_precision=True,
use_multi_tensor=True):
super().__init__()
self.learning_rate = learning_rate
self.momentum = momentum
self.weight_decay = weight_decay
self.grad_clip = grad_clip
self.multi_precision = multi_precision
self.use_multi_tensor = use_multi_tensor
def __call__(self, model_list):
# model_list is None in static graph
@ -54,6 +56,7 @@ class Momentum(object):
weight_decay=self.weight_decay,
grad_clip=self.grad_clip,
multi_precision=self.multi_precision,
use_multi_tensor=self.use_multi_tensor,
parameters=parameters)
return opt