27 lines
1.1 KiB
Python
27 lines
1.1 KiB
Python
from ml_lib.utils.config import Config
|
|
from models.conv_classifier import ConvClassifier
|
|
from models.bandwise_conv_classifier import BandwiseConvClassifier
|
|
from models.bandwise_conv_multihead_classifier import BandwiseConvMultiheadClassifier
|
|
from models.ensemble import Ensemble
|
|
from models.residual_conv_classifier import ResidualConvClassifier
|
|
from models.transformer_model import VisualTransformer
|
|
|
|
|
|
class MConfig(Config):
|
|
# TODO: There should be a way to automate this.
|
|
|
|
@property
|
|
def _model_map(self):
|
|
return dict(ConvClassifier=ConvClassifier,
|
|
CC=ConvClassifier,
|
|
BandwiseConvClassifier=BandwiseConvClassifier,
|
|
BCC=BandwiseConvClassifier,
|
|
BandwiseConvMultiheadClassifier=BandwiseConvMultiheadClassifier,
|
|
BCMC=BandwiseConvMultiheadClassifier,
|
|
Ensemble=Ensemble,
|
|
E=Ensemble,
|
|
ResidualConvClassifier=ResidualConvClassifier,
|
|
RCC=ResidualConvClassifier,
|
|
ViT=VisualTransformer
|
|
)
|