def setup_precision_plugin(self, plugin: PrecisionPlugin) -> None:
"""Attaches the precision plugin to the accelerator"""
model, optimizers, schedulers = plugin.connect(self.model,
self.optimizers,
self.lr_schedulers)
self.model = model
self.optimizers = optimizers
self.schedulers = schedulers
def __init__(
self,
precision_plugin: PrecisionPlugin = PrecisionPlugin(),
training_type_plugin: TrainingTypePlugin = SingleDevicePlugin(
torch.device(ipex.DEVICE)),
enable_bf16=False,
) -> None:
"""
Args:
precision_plugin: the plugin to handle precision-specific parts
training_type_plugin: the plugin to handle different training routines
"""
if enable_bf16:
# Automatically mix precision
ipex.enable_auto_mixed_precision(mixed_dtype=torch.bfloat16)
self.device = ipex.DEVICE
super().__init__(precision_plugin=precision_plugin,
training_type_plugin=training_type_plugin)
def precision_plugin(self) -> PrecisionPlugin:
return self._precision_plugin if self._precision_plugin is not None else PrecisionPlugin()