Registration Mechanism

The registration mechanism is an important part of HAT, which plays an important role in building config.

This section illustrates how to add a new module using the registration mechanism and how to use it in config by using a custom module as an example.

Module Customization

Here we use backbone as an example to demonstrate how to develop a new module (e.g., mobilenet).

Defining a New Backbone

Create a new file: hat/models/backbones/mobilenet.py.

import torch.nn as nn
from hat.registry import OBJECT_REGISTRY

__all__ = ["MobileNet"]

@OBJECT_REGISTRY.register
class MobileNet(nn.Module):
    def __init__(self, args1, args2):
        pass
    def forward(self, x):
        pass

Importing Newly Defined Module

You can import the newly defined module by adding the following line in hat/models/backbones/__init__.py:

from .mobilenet import MobileNet

Using the New Backbone in Config

model = dict(
    ...
    backbone=dict(
        type="MobileNet",
        args1=xxx,
        args2=xxx,
    )
    ...
)

Similarly, You can register and use other registrable module by using this method.