def __init__(self,
main_program,
inputs,
outputs,
parameters=None,
**kwargs):
super(PartialProgramLayer, self).__init__()
self._inputs = NestSequence(inputs)
self._outputs = NestSequence(outputs, need_check=True)
self._params = parameters if parameters is not None else []
self._build_strategy = kwargs.get('build_strategy', BuildStrategy())
assert isinstance(self._build_strategy, BuildStrategy)
self._origin_main_program = self._verify_program(main_program)
self._tmp_scope_vec = self._create_scope_vec()
self._cuda_graph_vec = self._create_cuda_graph_vec()
self._cuda_graph_capture_mode = ""
self._cuda_graph_pool_id = 0
# Set default mode to train
self.training = True
custom_white_list, custom_black_list = None, None
tracer = framework._dygraph_tracer()
if tracer:
custom_white_list, custom_black_list = tracer._get_amp_op_list()
# For AMP training
self._amp_list = AutoMixedPrecisionLists(
custom_white_list=custom_white_list,
custom_black_list=custom_black_list)
def sgd_optimizer(self, lr_policy, loss):
decayed_lr = self.get_lr(lr_policy)
optimizer = fluid.optimizer.Momentum(
learning_rate=decayed_lr,
momentum=self.momentum,
regularization=fluid.regularizer.L2Decay(
regularization_coeff=self.weight_decay),
)
if cfg.MODEL.FP16:
if cfg.MODEL.MODEL_NAME in ["pspnet"]:
custom_black_list = {"pool2d"}
else:
custom_black_list = {}
amp_lists = AutoMixedPrecisionLists(
custom_black_list=custom_black_list)
assert isinstance(cfg.MODEL.SCALE_LOSS, float) or isinstance(cfg.MODEL.SCALE_LOSS, str), \
"data type of MODEL.SCALE_LOSS must be float or str"
if isinstance(cfg.MODEL.SCALE_LOSS, float):
optimizer = decorate(optimizer,
amp_lists=amp_lists,
init_loss_scaling=cfg.MODEL.SCALE_LOSS,
use_dynamic_loss_scaling=False)
else:
assert cfg.MODEL.SCALE_LOSS.lower() in [
'dynamic'
], "if MODEL.SCALE_LOSS is a string,\
must be set as 'DYNAMIC'!"
optimizer = decorate(optimizer,
amp_lists=amp_lists,
use_dynamic_loss_scaling=True)
optimizer.minimize(loss)
return decayed_lr