def __call__(self,
global_step,
backbone: torch.nn.Module,
partial_fc: PartialFC = None):
if global_step > 100 and self.rank is 0:
torch.save(backbone.module.state_dict(),
os.path.join(self.output, "backbone.pth"))
if global_step > 100 and partial_fc is not None:
partial_fc.save_params()
def __call__(self,
global_step,
backbone: paddle.nn.Layer,
partial_fc: PartialFC = None):
if global_step > 100 and self.rank is 0:
paddle.save(
backbone.state_dict(),
os.path.join(self.output, self.model_name + ".pdparams"))
if global_step > 100 and partial_fc is not None:
partial_fc.save_params()