compressai_trainer.registry

Register models, runners, etc, to make them accessible via dynamic YAML configuration.

Registering maps a string name to a concrete creation method or class. This allows us to dynamically create an object depending on the given string name at runtime.

All models and runners should be registered and imported, as described in Defining a custom model and Defining a custom Runner training loop.

compressai/models/custom.py

from compressai.registry import register_model
from .base import CompressionModel

@register_model("my_custom_model")
class MyCustomModel(CompressionModel):
    def __init__(self, N, M):
        ...

compressai_trainer/runners/custom.py

from compressai.registry import register_runner
from .base import BaseRunner

@register_runner("CustomImageCompressionRunner")
class CustomImageCompressionRunner(BaseRunner):
    ...

torch

compressai_trainer.registry.torch.register_criterion(name: str)

Decorator for registering a criterion.

compressai_trainer.registry.torch.register_dataset(name: str)

Decorator for registering a dataset.

compressai_trainer.registry.torch.register_model(name: str)

Decorator for registering a model.

compressai_trainer.registry.torch.register_module(name: str)

Decorator for registering a module.

compressai_trainer.registry.torch.register_optimizer(name: str)

Decorator for registering a optimizer.

compressai_trainer.registry.torch.register_scheduler(name: str)

Decorator for registering a scheduler.

catalyst

compressai_trainer.registry.catalyst.register_callback(name: str)

Decorator for registering a callback.

compressai_trainer.registry.catalyst.register_runner(name: str)

Decorator for registering a runner.