diffengine.engine.hooks.compile_hook

Module Contents

Classes

CompileHook

Compile Hook.
class diffengine.engine.hooks.compile_hook.CompileHook(backend='inductor', mode=None, *, compile_main=False)[source]

Bases: mmengine.hooks.Hook

Compile Hook.

Args:

backend (str): The backend to use for compilation.

Defaults to “inductor”.

mode (str): The mode to use for compilation. Defaults to None. compile_main (bool): Whether to compile the main network like unet or

transformer. Defaults to False.

priority = VERY_LOW[source]
before_train(runner)[source]

Compile the model.

Args:

runner (Runner): The runner of the training process.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

Parameters:

  • backend (str) –

  • mode (str | None) –

  • compile_main (bool) –