init

lib/models/necks/__init__.py

@@ -0,0 +1,13 @@
+from .transformer_encoder import TransformerEncoder
+from .modality_completion import ModalityCompletion
+
+__all__ = ['TransformerEncoder', 'ModalityCompletion']
+
+type_mapping = {
+    'TransformerEncoder': TransformerEncoder,
+    'ModalityCompletion': ModalityCompletion
+}
+
+
+def build_neck(type, **kwargs):
+    return type_mapping[type](**kwargs)