Docker + smaller ConvNeXt-V2 models

.dockerignore

@@ -0,0 +1,8 @@
+*.swp
+**/__pycache__/**
+.idea/*
+ckpt/
+*.pth
+*.log
+*.txt
+.dockerignore

Dockerfile

@@ -0,0 +1,12 @@
+FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-devel
+
+WORKDIR /
+
+RUN  pip install timm==0.5.4
+
+COPY /pretrain/requirements.txt /
+RUN pip install --no-cache-dir -r  requirements.txt
+
+COPY . . 
+
+CMD ["bash"]

downstream_imagenet/arg.py

@@ -12,6 +12,11 @@
 
 HP_DEFAULT_NAMES = ['bs', 'ep', 'wp_ep', 'opt', 'base_lr', 'lr_scale', 'wd', 'mixup', 'rep_aug', 'drop_path', 'ema']
 HP_DEFAULT_VALUES = {
+    'convnext_atto':      (1024, 400, 20, 'adam', 0.0002,  0.7, 0.01, 0.8, 3, 0.3,  0.9999),
+    'convnext_femto':     (1024, 400, 20, 'adam', 0.0002,  0.7, 0.01, 0.8, 3, 0.3,  0.9999),
+    'convnext_pico':      (512,  400, 20, 'adam', 0.0002,  0.7, 0.01, 0.8, 3, 0.3,  0.9999),
+    'convnext_nano':      (512,  400, 20, 'adam', 0.0002,  0.7, 0.01, 0.8, 3, 0.3,  0.9999),
+    'convnext_tiny':      (256,  400, 20, 'adam', 0.0002,  0.7, 0.01, 0.8, 3, 0.3,  0.9999),
     'convnext_small':     (4096, 400, 20, 'adam', 0.0002,  0.7, 0.01, 0.8, 3, 0.3,  0.9999),
     'convnext_base':      (4096, 400, 20, 'adam', 0.0001,  0.7, 0.01, 0.8, 3, 0.4,  0.9999),
     'convnext_large':     (4096, 200, 10, 'adam', 0.0001,  0.7, 0.02, 0.8, 3, 0.5,  0.9999),

downstream_imagenet/models/convnext_official.py

@@ -154,6 +154,42 @@ def forward(self, x):
     "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
 }
 
+@register_model
+def convnext_atto(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 6, 2], dims=[40, 80, 160, 320], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_atto_22k'] if in_22k else model_urls['convnext_atto_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_femto(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 6, 2], dims=[48, 96, 192, 384], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_femto_22k'] if in_22k else model_urls['convnext_femto_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_pico(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 6, 2], dims=[64, 128, 256, 512], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_pico_22k'] if in_22k else model_urls['convnext_pico_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
+@register_model
+def convnext_nano(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 8, 2], dims=[80, 160, 320, 640], **kwargs)
+    if pretrained:
+        url = model_urls['convnext_nano_22k'] if in_22k else model_urls['convnext_nano_1k']
+        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model.load_state_dict(checkpoint["model"])
+    return model
+
 @register_model
 def convnext_tiny(pretrained=False,in_22k=False, **kwargs):
     model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)

pretrain/models/convnext.py

@@ -99,12 +99,29 @@ def get_classifier(self):
     def extra_repr(self):
         return f'drop_path_rate={self.drop_path_rate}, layer_scale_init_value={self.layer_scale_init_value:g}'
 
+@register_model
+def convnext_atto(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 6, 2], dims=[40, 80, 160, 320], **kwargs)
+    return model
 
 @register_model
-def convnext_tiny(pretrained=False, in_22k=False, **kwargs):
-    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+def convnext_femto(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 6, 2], dims=[48, 96, 192, 384], **kwargs)
     return model
 
+@register_model
+def convnext_pico(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 6, 2], dims=[64, 128, 256, 512], **kwargs)
+
+@register_model
+def convnext_nano(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[2, 2, 8, 2], dims=[80, 160, 320, 640], **kwargs)
+    return model
+
+@register_model
+def convnext_tiny(pretrained=False,in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    return model
 
 @register_model
 def convnext_small(pretrained=False, in_22k=False, **kwargs):

pretrain/requirements.txt

@@ -4,3 +4,4 @@ Pillow
 typed-argument-parser
 timm==0.5.4
 tensorboardx
+tensorboard