Add files using upload-large-folder tool

CION_ReIDZoo/CION_Finetune/FastReID/configs/Market1501/AGW_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: "../Base-AGW.yml"
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: "logs/market1501/agw_R50"

CION_ReIDZoo/CION_Finetune/FastReID/configs/VeRi/sbs_R50-ibn.yml

@@ -0,0 +1,32 @@
+_BASE_: "../Base-Strongerbaseline.yml"
+
+INPUT:
+  SIZE_TRAIN: [256, 256]
+  SIZE_TEST: [256, 256]
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+
+SOLVER:
+  OPT: "SGD"
+  BASE_LR: 0.01
+  ETA_MIN_LR: 7.7e-5
+
+  IMS_PER_BATCH: 64
+  MAX_ITER: 60
+  DELAY_ITERS: 30
+  WARMUP_ITERS: 10
+  FREEZE_ITERS: 10
+
+  CHECKPOINT_PERIOD: 20
+
+DATASETS:
+  NAMES: ("VeRi",)
+  TESTS: ("VeRi",)
+
+TEST:
+  EVAL_PERIOD: 20
+  IMS_PER_BATCH: 128
+
+OUTPUT_DIR: "logs/veri/sbs_R50-ibn"

CION_ReIDZoo/CION_Finetune/FastReID/configs/VehicleID/bagtricks_R50-ibn.yml

@@ -0,0 +1,35 @@
+_BASE_: "../Base-bagtricks.yml"
+
+INPUT:
+  SIZE_TRAIN: [256, 256]
+  SIZE_TEST: [256, 256]
+
+MODEL:
+  BACKBONE:
+    WITH_IBN: True
+  HEADS:
+    POOL_LAYER: gempool
+  LOSSES:
+    TRI:
+      HARD_MINING: False
+      MARGIN: 0.0
+
+DATASETS:
+  NAMES: ("VehicleID",)
+  TESTS: ("SmallVehicleID", "MediumVehicleID", "LargeVehicleID",)
+
+SOLVER:
+  BIAS_LR_FACTOR: 1.
+
+  IMS_PER_BATCH: 512
+  MAX_ITER: 60
+  STEPS: [30, 50]
+  WARMUP_ITERS: 10
+
+  CHECKPOINT_PERIOD: 20
+
+TEST:
+  EVAL_PERIOD: 20
+  IMS_PER_BATCH: 128
+
+OUTPUT_DIR: "logs/vehicleid/bagtricks_R50-ibn_4gpu"

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/Caffe/ReadMe.md

@@ -0,0 +1,21 @@
+# The Caffe in nn_tools Provides some convenient API
+If there are some problem in parse your prototxt or caffemodel, Please replace
+the caffe.proto with your own version and compile it with command
+                   `protoc --python_out ./ caffe.proto`
+
+## caffe_net.py
+Using `from nn_tools.Caffe import caffe_net` to import this model
+### Prototxt
++ `net=caffe_net.Prototxt(file_name)` to open a prototxt file
++ `net.init_caffemodel(caffe_cmd_path='caffe')` to generate a caffemodel file in the current work directory \
+if your `caffe` cmd not in the $PATH, specify your caffe cmd path by the `caffe_cmd_path` kwargs.
+### Caffemodel
++ `net=caffe_net.Caffemodel(file_name)` to open a caffemodel
++ `net.save_prototxt(path)` to save the caffemodel to a prototxt file (not containing the weight data)
++ `net.get_layer_data(layer_name)` return the numpy ndarray data of the layer
++ `net.set_layer_date(layer_name, datas)` specify the data of one layer in the caffemodel .`datas` is normally a list of numpy ndarray `[weights,bias]`
++ `net.save(path)` save the changed caffemodel
+### Functions for both Prototxt and Caffemodel
++ `net.add_layer(layer_params,before='',after='')` add a new layer with `Layer_Param` object
++ `net.remove_layer_by_name(layer_name)` 
++ `net.get_layer_by_name(layer_name)` or `net.layer(layer_name)` get the raw Layer object defined in caffe_pb2

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/Caffe/caffe_net.py

@@ -0,0 +1,139 @@
+from __future__ import absolute_import
+from . import caffe_pb2 as pb
+import google.protobuf.text_format as text_format
+import numpy as np
+from .layer_param import Layer_param
+
+class _Net(object):
+    def __init__(self):
+        self.net=pb.NetParameter()
+
+    def layer_index(self,layer_name):
+        # find a layer's index by name. if the layer was found, return the layer position in the net, else return -1.
+        for i, layer in enumerate(self.net.layer):
+            if layer.name == layer_name:
+                return i
+
+    def add_layer(self,layer_params,before='',after=''):
+        # find the before of after layer's position
+        index = -1
+        if after != '':
+            index = self.layer_index(after) + 1
+        if before != '':
+            index = self.layer_index(before)
+        new_layer = pb.LayerParameter()
+        new_layer.CopyFrom(layer_params.param)
+        #insert the layer into the layer protolist
+        if index != -1:
+            self.net.layer.add()
+            for i in range(len(self.net.layer) - 1, index, -1):
+                self.net.layer[i].CopyFrom(self.net.layer[i - 1])
+            self.net.layer[index].CopyFrom(new_layer)
+        else:
+            self.net.layer.extend([new_layer])
+
+    def remove_layer_by_name(self,layer_name):
+        for i,layer in enumerate(self.net.layer):
+            if layer.name == layer_name:
+                del self.net.layer[i]
+                return
+        raise(AttributeError, "cannot found layer %s" % str(layer_name))
+
+    def get_layer_by_name(self, layer_name):
+        # get the layer by layer_name
+        for layer in self.net.layer:
+            if layer.name == layer_name:
+                return layer
+        raise(AttributeError, "cannot found layer %s" % str(layer_name))
+
+    def save_prototxt(self,path):
+        prototxt=pb.NetParameter()
+        prototxt.CopyFrom(self.net)
+        for layer in prototxt.layer:
+            del layer.blobs[:]
+        with open(path,'w') as f:
+            f.write(text_format.MessageToString(prototxt))
+
+    def layer(self,layer_name):
+        return self.get_layer_by_name(layer_name)
+
+    def layers(self):
+        return list(self.net.layer)
+
+
+
+class Prototxt(_Net):
+    def __init__(self,file_name=''):
+        super(Prototxt,self).__init__()
+        self.file_name=file_name
+        if file_name!='':
+            f = open(file_name,'r')
+            text_format.Parse(f.read(), self.net)
+            pass
+
+    def init_caffemodel(self,caffe_cmd_path='caffe'):
+        """
+        :param caffe_cmd_path: The shell command of caffe, normally at <path-to-caffe>/build/tools/caffe
+        """
+        s=pb.SolverParameter()
+        s.train_net=self.file_name
+        s.max_iter=0
+        s.base_lr=1
+        s.solver_mode = pb.SolverParameter.CPU
+        s.snapshot_prefix='./nn'
+        with open('/tmp/nn_tools_solver.prototxt','w') as f:
+            f.write(str(s))
+        import os
+        os.system('%s train --solver /tmp/nn_tools_solver.prototxt'%caffe_cmd_path)
+
+class Caffemodel(_Net):
+    def __init__(self, file_name=''):
+        super(Caffemodel,self).__init__()
+        # caffe_model dir
+        if file_name!='':
+            f = open(file_name,'rb')
+            self.net.ParseFromString(f.read())
+            f.close()
+
+    def save(self, path):
+        with open(path,'wb') as f:
+            f.write(self.net.SerializeToString())
+
+    def add_layer_with_data(self,layer_params,datas, before='', after=''):
+        """
+        Args:
+            layer_params:A Layer_Param object
+            datas:a fixed dimension numpy object list
+            after: put the layer after a specified layer
+            before: put the layer before a specified layer
+        """
+        self.add_layer(layer_params,before,after)
+        new_layer =self.layer(layer_params.name)
+
+        #process blobs
+        del new_layer.blobs[:]
+        for data in datas:
+            new_blob=new_layer.blobs.add()
+            for dim in data.shape:
+                new_blob.shape.dim.append(dim)
+            new_blob.data.extend(data.flatten().astype(float))
+
+    def get_layer_data(self,layer_name):
+        layer=self.layer(layer_name)
+        datas=[]
+        for blob in layer.blobs:
+            shape=list(blob.shape.dim)
+            data=np.array(blob.data).reshape(shape)
+            datas.append(data)
+        return datas
+
+    def set_layer_data(self,layer_name,datas):
+        # datas is normally a list of [weights,bias]
+        layer=self.layer(layer_name)
+        for blob,data in zip(layer.blobs,datas):
+            blob.data[:]=data.flatten()
+            pass
+
+class Net():
+    def __init__(self,*args,**kwargs):
+        raise(TypeError,'the class Net is no longer used, please use Caffemodel or Prototxt instead')

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/Caffe/layer_param.py

@@ -0,0 +1,158 @@
+from __future__ import absolute_import
+from . import caffe_pb2 as pb
+import numpy as np
+
+def pair_process(item,strict_one=True):
+    if hasattr(item,'__iter__'):
+        for i in item:
+            if i!=item[0]:
+                if strict_one:
+                    raise ValueError("number in item {} must be the same".format(item))
+                else:
+                    print("IMPORTANT WARNING: number in item {} must be the same".format(item))
+        return item[0]
+    return item
+
+def pair_reduce(item):
+    if hasattr(item,'__iter__'):
+        for i in item:
+            if i!=item[0]:
+                return item
+        return [item[0]]
+    return [item]
+
+class Layer_param():
+    def __init__(self,name='',type='',top=(),bottom=()):
+        self.param=pb.LayerParameter()
+        self.name=self.param.name=name
+        self.type=self.param.type=type
+
+        self.top=self.param.top
+        self.top.extend(top)
+        self.bottom=self.param.bottom
+        self.bottom.extend(bottom)
+
+    def fc_param(self, num_output, weight_filler='xavier', bias_filler='constant',has_bias=True):
+        if self.type != 'InnerProduct':
+            raise TypeError('the layer type must be InnerProduct if you want set fc param')
+        fc_param = pb.InnerProductParameter()
+        fc_param.num_output = num_output
+        fc_param.weight_filler.type = weight_filler
+        fc_param.bias_term = has_bias
+        if has_bias:
+            fc_param.bias_filler.type = bias_filler
+        self.param.inner_product_param.CopyFrom(fc_param)
+
+    def conv_param(self, num_output, kernel_size, stride=(1), pad=(0,),
+                   weight_filler_type='xavier', bias_filler_type='constant',
+                   bias_term=True, dilation=None,groups=None):
+        """
+        add a conv_param layer if you spec the layer type "Convolution"
+        Args:
+            num_output: a int
+            kernel_size: int list
+            stride: a int list
+            weight_filler_type: the weight filer type
+            bias_filler_type: the bias filler type
+        Returns:
+        """
+        if self.type not in ['Convolution','Deconvolution']:
+            raise TypeError('the layer type must be Convolution or Deconvolution if you want set conv param')
+        conv_param=pb.ConvolutionParameter()
+        conv_param.num_output=num_output
+        conv_param.kernel_size.extend(pair_reduce(kernel_size))
+        conv_param.stride.extend(pair_reduce(stride))
+        conv_param.pad.extend(pair_reduce(pad))
+        conv_param.bias_term=bias_term
+        conv_param.weight_filler.type=weight_filler_type
+        if bias_term:
+            conv_param.bias_filler.type = bias_filler_type
+        if dilation:
+            conv_param.dilation.extend(pair_reduce(dilation))
+        if groups:
+            conv_param.group=groups
+        self.param.convolution_param.CopyFrom(conv_param)
+
+    def pool_param(self,type='MAX',kernel_size=2,stride=2,pad=None, ceil_mode = False):
+        pool_param=pb.PoolingParameter()
+        pool_param.pool=pool_param.PoolMethod.Value(type)
+        pool_param.kernel_size=pair_process(kernel_size)
+        pool_param.stride=pair_process(stride)
+        pool_param.ceil_mode=ceil_mode
+        if pad:
+            if isinstance(pad,tuple):
+                pool_param.pad_h = pad[0]
+                pool_param.pad_w = pad[1]
+            else:
+                pool_param.pad=pad
+        self.param.pooling_param.CopyFrom(pool_param)
+
+    def batch_norm_param(self,use_global_stats=0,moving_average_fraction=None,eps=None):
+        bn_param=pb.BatchNormParameter()
+        bn_param.use_global_stats=use_global_stats
+        if moving_average_fraction:
+            bn_param.moving_average_fraction=moving_average_fraction
+        if eps:
+            bn_param.eps = eps
+        self.param.batch_norm_param.CopyFrom(bn_param)
+
+    # layer
+    # {
+    #     name: "upsample_layer"
+    #     type: "Upsample"
+    #     bottom: "some_input_feature_map"
+    #     bottom: "some_input_pool_index"
+    #     top: "some_output"
+    #     upsample_param {
+    #         upsample_h: 224
+    #         upsample_w: 224
+    #     }
+    # }
+    def upsample_param(self,size=None, scale_factor=None):
+        upsample_param=pb.UpsampleParameter()
+        if scale_factor:
+            if isinstance(scale_factor,int):
+                upsample_param.scale = scale_factor
+            else:
+                upsample_param.scale_h = scale_factor[0]
+                upsample_param.scale_w = scale_factor[1]
+
+        if size:
+            if isinstance(size,int):
+                upsample_param.upsample_h = size
+            else:
+                upsample_param.upsample_h = size[0]
+                upsample_param.upsample_w = size[1]
+                #upsample_param.upsample_h = size[0] * scale_factor
+                #upsample_param.upsample_w = size[1] * scale_factor
+        self.param.upsample_param.CopyFrom(upsample_param)
+    def interp_param(self,size=None, scale_factor=None):
+        interp_param=pb.InterpParameter()
+        if scale_factor:
+            if isinstance(scale_factor,int):
+                interp_param.zoom_factor = scale_factor
+
+        if size:
+            print('size:', size)
+            interp_param.height = size[0]
+            interp_param.width = size[1]
+        self.param.interp_param.CopyFrom(interp_param)
+
+    def add_data(self,*args):
+        """Args are data numpy array
+        """
+        del self.param.blobs[:]
+        for data in args:
+            new_blob = self.param.blobs.add()
+            for dim in data.shape:
+                new_blob.shape.dim.append(dim)
+            new_blob.data.extend(data.flatten().astype(float))
+
+    def set_params_by_dict(self,dic):
+        pass
+
+    def copy_from(self,layer_param):
+        pass
+
+def set_enum(param,key,value):
+    setattr(param,key,param.Value(value))

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/README.md

@@ -0,0 +1,160 @@
+# Model Deployment
+
+This directory contains:
+
+1. The scripts that convert a fastreid model to Caffe/ONNX/TRT format.
+
+2. The exmpales that load a R50 baseline model in Caffe/ONNX/TRT and run inference.
+
+## Tutorial
+
+### Caffe Convert
+
+<details>
+<summary>step-to-step pipeline for caffe convert</summary>
+
+This is a tiny example for converting fastreid-baseline in `meta_arch` to Caffe model, if you want to convert more complex architecture, you need to customize more things.
+
+1. Run `caffe_export.py` to get the converted Caffe model,
+
+    ```bash
+    python caffe_export.py --config-file root-path/market1501/bagtricks_R50/config.yml --name "baseline_R50" --output outputs/caffe_model --opts MODEL.WEIGHTS root-path/logs/market1501/bagtricks_R50/model_final.pth
+    ```
+
+    then you can check the Caffe model and prototxt in `outputs/caffe_model`.
+
+2. Change `prototxt` following next three steps:
+
+   1) Edit `max_pooling` in `baseline_R50.prototxt` like this
+
+        ```prototxt
+        layer {
+            name: "max_pool1"
+            type: "Pooling"
+            bottom: "relu_blob1"
+            top: "max_pool_blob1"
+            pooling_param {
+                pool: MAX
+                kernel_size: 3
+                stride: 2
+                pad: 0 # 1
+                # ceil_mode: false
+            }
+        }
+        ```
+
+   2) Add `avg_pooling` right place in `baseline_R50.prototxt`
+
+        ```prototxt
+        layer {
+            name: "avgpool1"
+            type: "Pooling"
+            bottom: "relu_blob49"
+            top: "avgpool_blob1"
+            pooling_param {
+                pool: AVE
+                global_pooling: true
+            }
+        }
+        ```
+
+   3) Change the last layer `top` name to `output`
+
+        ```prototxt
+        layer {
+            name: "bn_scale54"
+            type: "Scale"
+            bottom: "batch_norm_blob54"
+            top: "output" # bn_norm_blob54
+            scale_param {
+                bias_term: true
+            }
+        }
+        ```
+
+3. (optional) You can open [Netscope](https://ethereon.github.io/netscope/quickstart.html), then enter you network `prototxt` to visualize the network.
+
+4. Run `caffe_inference.py` to save Caffe model features with input images
+
+   ```bash
+    python caffe_inference.py --model-def outputs/caffe_model/baseline_R50.prototxt \
+    --model-weights outputs/caffe_model/baseline_R50.caffemodel \
+    --input test_data/*.jpg --output caffe_output
+   ```
+
+5. Run `demo/demo.py` to get fastreid model features with the same input images, then verify that Caffe and PyTorch are computing the same value for the network.
+
+    ```python
+    np.testing.assert_allclose(torch_out, ort_out, rtol=1e-3, atol=1e-6)
+    ```
+
+</details>
+
+### ONNX Convert
+
+<details>
+<summary>step-to-step pipeline for onnx convert</summary>
+
+This is a tiny example for converting fastreid-baseline in `meta_arch` to ONNX model. ONNX supports most operators in pytorch as far as I know and if some operators are not supported by ONNX, you need to customize these.
+
+1. Run `onnx_export.py` to get the converted ONNX model,
+
+    ```bash
+    python onnx_export.py --config-file root-path/bagtricks_R50/config.yml --name "baseline_R50" --output outputs/onnx_model --opts MODEL.WEIGHTS root-path/logs/market1501/bagtricks_R50/model_final.pth
+    ```
+
+    then you can check the ONNX model in `outputs/onnx_model`.
+
+2. (optional) You can use [Netron](https://github.com/lutzroeder/netron) to visualize the network.
+
+3. Run `onnx_inference.py` to save ONNX model features with input images
+
+   ```bash
+    python onnx_inference.py --model-path outputs/onnx_model/baseline_R50.onnx \
+    --input test_data/*.jpg --output onnx_output
+   ```
+
+4. Run `demo/demo.py` to get fastreid model features with the same input images, then verify that ONNX Runtime and PyTorch are computing the same value for the network.
+
+    ```python
+    np.testing.assert_allclose(torch_out, ort_out, rtol=1e-3, atol=1e-6)
+    ```
+
+</details>
+
+### TensorRT Convert
+
+<details>
+<summary>step-to-step pipeline for trt convert</summary>
+
+This is a tiny example for converting fastreid-baseline in `meta_arch` to TRT model. We use [tiny-tensorrt](https://github.com/zerollzeng/tiny-tensorrt), which is a simple and easy-to-use nvidia TensorRT warpper, to get the model converted to tensorRT.
+
+First you need to convert the pytorch model to ONNX format following [ONNX Convert](https://github.com/JDAI-CV/fast-reid/tree/master/tools/deploy#onnx-convert), and you need to remember your `output` name. Then you can convert ONNX model to TensorRT following instructions below.
+
+1. Run command line below to get the converted TRT model from ONNX model,
+
+    ```bash
+
+    python trt_export.py --name "baseline_R50" --output outputs/trt_model --onnx-model outputs/onnx_model/baseline.onnx --heighi 256 --width 128
+    ```
+
+    then you can check the TRT model in `outputs/trt_model`.
+
+2. Run `trt_inference.py` to save TRT model features with input images
+
+   ```bash
+    python onnx_inference.py --model-path outputs/trt_model/baseline.engine \
+    --input test_data/*.jpg --output trt_output --output-name trt_model_outputname
+   ```
+
+3. Run `demo/demo.py` to get fastreid model features with the same input images, then verify that TensorRT and PyTorch are computing the same value for the network.
+
+    ```python
+    np.testing.assert_allclose(torch_out, ort_out, rtol=1e-3, atol=1e-6)
+    ```
+
+</details>
+
+## Acknowledgements
+
+Thank to [CPFLAME](https://github.com/CPFLAME), [gcong18](https://github.com/gcong18), [YuxiangJohn](https://github.com/YuxiangJohn) and [wiggin66](https://github.com/wiggin66) at JDAI Model Acceleration Group for help in PyTorch model converting.

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/caffe_export.py

@@ -0,0 +1,78 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: [email protected]
+"""
+
+import argparse
+
+import torch
+import sys
+sys.path.append('../../')
+
+import pytorch_to_caffe
+from fastreid.config import get_cfg
+from fastreid.modeling.meta_arch import build_model
+from fastreid.utils.file_io import PathManager
+from fastreid.utils.checkpoint import Checkpointer
+from fastreid.utils.logger import setup_logger
+
+logger = setup_logger(name='caffe_export')
+
+
+def setup_cfg(args):
+    cfg = get_cfg()
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+    return cfg
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Convert Pytorch to Caffe model")
+
+    parser.add_argument(
+        "--config-file",
+        metavar="FILE",
+        help="path to config file",
+    )
+    parser.add_argument(
+        "--name",
+        default="baseline",
+        help="name for converted model"
+    )
+    parser.add_argument(
+        "--output",
+        default='caffe_model',
+        help='path to save converted caffe model'
+    )
+    parser.add_argument(
+        "--opts",
+        help="Modify config options using the command-line 'KEY VALUE' pairs",
+        default=[],
+        nargs=argparse.REMAINDER,
+    )
+    return parser
+
+
+if __name__ == '__main__':
+    args = get_parser().parse_args()
+    cfg = setup_cfg(args)
+
+    cfg.defrost()
+    cfg.MODEL.BACKBONE.PRETRAIN = False
+    cfg.MODEL.HEADS.POOL_LAYER = "identity"
+    cfg.MODEL.BACKBONE.WITH_NL = False
+
+    model = build_model(cfg)
+    Checkpointer(model).load(cfg.MODEL.WEIGHTS)
+    model.eval()
+    logger.info(model)
+
+    inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1]).to(torch.device(cfg.MODEL.DEVICE))
+    PathManager.mkdirs(args.output)
+    pytorch_to_caffe.trans_net(model, inputs, args.name)
+    pytorch_to_caffe.save_prototxt(f"{args.output}/{args.name}.prototxt")
+    pytorch_to_caffe.save_caffemodel(f"{args.output}/{args.name}.caffemodel")
+
+    logger.info(f"Export caffe model in {args.output} sucessfully!")

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/caffe_inference.py

@@ -0,0 +1,95 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: [email protected]
+"""
+
+import caffe
+import tqdm
+import glob
+import os
+import cv2
+import numpy as np
+
+caffe.set_mode_gpu()
+
+import argparse
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Caffe model inference")
+
+    parser.add_argument(
+        "--model-def",
+        default="logs/test_caffe/baseline_R50.prototxt",
+        help="caffe model prototxt"
+    )
+    parser.add_argument(
+        "--model-weights",
+        default="logs/test_caffe/baseline_R50.caffemodel",
+        help="caffe model weights"
+    )
+    parser.add_argument(
+        "--input",
+        nargs="+",
+        help="A list of space separated input images; "
+             "or a single glob pattern such as 'directory/*.jpg'",
+    )
+    parser.add_argument(
+        "--output",
+        default='caffe_output',
+        help='path to save converted caffe model'
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=256,
+        help="height of image"
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=128,
+        help="width of image"
+    )
+    return parser
+
+
+def preprocess(image_path, image_height, image_width):
+    original_image = cv2.imread(image_path)
+    # the model expects RGB inputs
+    original_image = original_image[:, :, ::-1]
+
+    # Apply pre-processing to image.
+    image = cv2.resize(original_image, (image_width, image_height), interpolation=cv2.INTER_CUBIC)
+    image = image.astype("float32").transpose(2, 0, 1)[np.newaxis]  # (1, 3, h, w)
+    image = (image - np.array([0.485 * 255, 0.456 * 255, 0.406 * 255]).reshape((1, -1, 1, 1))) / np.array(
+        [0.229 * 255, 0.224 * 255, 0.225 * 255]).reshape((1, -1, 1, 1))
+    return image
+
+
+def normalize(nparray, order=2, axis=-1):
+    """Normalize a N-D numpy array along the specified axis."""
+    norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
+    return nparray / (norm + np.finfo(np.float32).eps)
+
+
+if __name__ == "__main__":
+    args = get_parser().parse_args()
+
+    net = caffe.Net(args.model_def, args.model_weights, caffe.TEST)
+    net.blobs['blob1'].reshape(1, 3, args.height, args.width)
+
+    if not os.path.exists(args.output): os.makedirs(args.output)
+
+    if args.input:
+        if os.path.isdir(args.input[0]):
+            args.input = glob.glob(os.path.expanduser(args.input[0]))
+            assert args.input, "The input path(s) was not found"
+        for path in tqdm.tqdm(args.input):
+            image = preprocess(path, args.height, args.width)
+            net.blobs['blob1'].data[...] = image
+            feat = net.forward()['output']
+            feat = normalize(feat[..., 0, 0], axis=1)
+            np.save(os.path.join(args.output, path.replace('.jpg', '.npy').split('/')[-1]), feat)
+

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/onnx_export.py

@@ -0,0 +1,146 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: [email protected]
+"""
+
+import argparse
+import io
+import sys
+
+import onnx
+import torch
+from onnxsim import simplify
+from torch.onnx import OperatorExportTypes
+
+sys.path.append('../../')
+
+from fastreid.config import get_cfg
+from fastreid.modeling.meta_arch import build_model
+from fastreid.utils.file_io import PathManager
+from fastreid.utils.checkpoint import Checkpointer
+from fastreid.utils.logger import setup_logger
+
+logger = setup_logger(name='onnx_export')
+
+
+def setup_cfg(args):
+    cfg = get_cfg()
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+    return cfg
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Convert Pytorch to ONNX model")
+
+    parser.add_argument(
+        "--config-file",
+        metavar="FILE",
+        help="path to config file",
+    )
+    parser.add_argument(
+        "--name",
+        default="baseline",
+        help="name for converted model"
+    )
+    parser.add_argument(
+        "--output",
+        default='onnx_model',
+        help='path to save converted onnx model'
+    )
+    parser.add_argument(
+        "--opts",
+        help="Modify config options using the command-line 'KEY VALUE' pairs",
+        default=[],
+        nargs=argparse.REMAINDER,
+    )
+    return parser
+
+
+def remove_initializer_from_input(model):
+    if model.ir_version < 4:
+        print(
+            'Model with ir_version below 4 requires to include initilizer in graph input'
+        )
+        return
+
+    inputs = model.graph.input
+    name_to_input = {}
+    for input in inputs:
+        name_to_input[input.name] = input
+
+    for initializer in model.graph.initializer:
+        if initializer.name in name_to_input:
+            inputs.remove(name_to_input[initializer.name])
+
+    return model
+
+
+def export_onnx_model(model, inputs):
+    """
+    Trace and export a model to onnx format.
+    Args:
+        model (nn.Module):
+        inputs (torch.Tensor): the model will be called by `model(*inputs)`
+    Returns:
+        an onnx model
+    """
+    assert isinstance(model, torch.nn.Module)
+
+    # make sure all modules are in eval mode, onnx may change the training state
+    # of the module if the states are not consistent
+    def _check_eval(module):
+        assert not module.training
+
+    model.apply(_check_eval)
+
+    # Export the model to ONNX
+    with torch.no_grad():
+        with io.BytesIO() as f:
+            torch.onnx.export(
+                model,
+                inputs,
+                f,
+                operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,
+                # verbose=True,  # NOTE: uncomment this for debugging
+                # export_params=True,
+            )
+            onnx_model = onnx.load_from_string(f.getvalue())
+
+    # Apply ONNX's Optimization
+    all_passes = onnx.optimizer.get_available_passes()
+    passes = ["extract_constant_to_initializer", "eliminate_unused_initializer", "fuse_bn_into_conv"]
+    assert all(p in all_passes for p in passes)
+    onnx_model = onnx.optimizer.optimize(onnx_model, passes)
+    return onnx_model
+
+
+if __name__ == '__main__':
+    args = get_parser().parse_args()
+    cfg = setup_cfg(args)
+
+    cfg.defrost()
+    cfg.MODEL.BACKBONE.PRETRAIN = False
+    if cfg.MODEL.HEADS.POOL_LAYER == 'fastavgpool':
+        cfg.MODEL.HEADS.POOL_LAYER = 'avgpool'
+    model = build_model(cfg)
+    Checkpointer(model).load(cfg.MODEL.WEIGHTS)
+    model.eval()
+    logger.info(model)
+
+    inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1])
+    onnx_model = export_onnx_model(model, inputs)
+
+    model_simp, check = simplify(onnx_model)
+
+    model_simp = remove_initializer_from_input(model_simp)
+
+    assert check, "Simplified ONNX model could not be validated"
+
+    PathManager.mkdirs(args.output)
+
+    onnx.save_model(model_simp, f"{args.output}/{args.name}.onnx")
+
+    logger.info(f"Export onnx model in {args.output} successfully!")

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/onnx_inference.py

@@ -0,0 +1,85 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: [email protected]
+"""
+
+import argparse
+import glob
+import os
+
+import cv2
+import numpy as np
+import onnxruntime
+import tqdm
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="onnx model inference")
+
+    parser.add_argument(
+        "--model-path",
+        default="onnx_model/baseline.onnx",
+        help="onnx model path"
+    )
+    parser.add_argument(
+        "--input",
+        nargs="+",
+        help="A list of space separated input images; "
+             "or a single glob pattern such as 'directory/*.jpg'",
+    )
+    parser.add_argument(
+        "--output",
+        default='onnx_output',
+        help='path to save converted caffe model'
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=256,
+        help="height of image"
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=128,
+        help="width of image"
+    )
+    return parser
+
+
+def preprocess(image_path, image_height, image_width):
+    original_image = cv2.imread(image_path)
+    # the model expects RGB inputs
+    original_image = original_image[:, :, ::-1]
+
+    # Apply pre-processing to image.
+    img = cv2.resize(original_image, (image_width, image_height), interpolation=cv2.INTER_CUBIC)
+    img = img.astype("float32").transpose(2, 0, 1)[np.newaxis]  # (1, 3, h, w)
+    return img
+
+
+def normalize(nparray, order=2, axis=-1):
+    """Normalize a N-D numpy array along the specified axis."""
+    norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
+    return nparray / (norm + np.finfo(np.float32).eps)
+
+
+if __name__ == "__main__":
+    args = get_parser().parse_args()
+
+    ort_sess = onnxruntime.InferenceSession(args.model_path)
+
+    input_name = ort_sess.get_inputs()[0].name
+
+    if not os.path.exists(args.output): os.makedirs(args.output)
+
+    if args.input:
+        if os.path.isdir(args.input[0]):
+            args.input = glob.glob(os.path.expanduser(args.input[0]))
+            assert args.input, "The input path(s) was not found"
+        for path in tqdm.tqdm(args.input):
+            image = preprocess(path, args.height, args.width)
+            feat = ort_sess.run(None, {input_name: image})[0]
+            feat = normalize(feat, axis=1)
+            np.save(os.path.join(args.output, path.replace('.jpg', '.npy').split('/')[-1]), feat)

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/pytorch_to_caffe.py

@@ -0,0 +1,747 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+import traceback
+from Caffe import caffe_net
+import torch.nn.functional as F
+from torch.autograd import Variable
+from Caffe import layer_param
+from torch.nn.modules.utils import _pair
+import numpy as np
+import math
+from torch.nn.modules.utils import _list_with_default
+
+"""
+How to support a new layer type:
+ layer_name=log.add_layer(layer_type_name)
+ top_blobs=log.add_blobs(<output of that layer>)
+ layer=caffe_net.Layer_param(xxx)
+ <set layer parameters>
+ [<layer.add_data(*datas)>]
+ log.cnet.add_layer(layer)
+
+Please MUTE the inplace operations to avoid not find in graph
+"""
+
+
+# TODO: support the inplace output of the layers
+
+class Blob_LOG():
+    def __init__(self):
+        self.data = {}
+
+    def __setitem__(self, key, value):
+        self.data[key] = value
+
+    def __getitem__(self, key):
+        return self.data[key]
+
+    def __len__(self):
+        return len(self.data)
+
+
+NET_INITTED = False
+
+
+# 转换原理解析：通过记录
+class TransLog(object):
+    def __init__(self):
+        """
+        doing init() with inputs Variable before using it
+        """
+        self.layers = {}
+        self.detail_layers = {}
+        self.detail_blobs = {}
+        self._blobs = Blob_LOG()
+        self._blobs_data = []
+        self.cnet = caffe_net.Caffemodel('')
+        self.debug = True
+
+    def init(self, inputs):
+        """
+        :param inputs: is a list of input variables
+        """
+        self.add_blobs(inputs)
+
+    def add_layer(self, name='layer'):
+        if name in self.layers:
+            return self.layers[name]
+        if name not in self.detail_layers.keys():
+            self.detail_layers[name] = 0
+        self.detail_layers[name] += 1
+        name = '{}{}'.format(name, self.detail_layers[name])
+        self.layers[name] = name
+        if self.debug:
+            print("{} was added to layers".format(self.layers[name]))
+        return self.layers[name]
+
+    def add_blobs(self, blobs, name='blob', with_num=True):
+        rst = []
+        for blob in blobs:
+            self._blobs_data.append(blob)  # to block the memory address be rewrited
+            blob_id = int(id(blob))
+            if name not in self.detail_blobs.keys():
+                self.detail_blobs[name] = 0
+            self.detail_blobs[name] += 1
+            if with_num:
+                rst.append('{}{}'.format(name, self.detail_blobs[name]))
+            else:
+                rst.append('{}'.format(name))
+            if self.debug:
+                print("{}:{} was added to blobs".format(blob_id, rst[-1]))
+            print('Add blob {} : {}'.format(rst[-1].center(21), blob.size()))
+            self._blobs[blob_id] = rst[-1]
+        return rst
+
+    def blobs(self, var):
+        var = id(var)
+        if self.debug:
+            print("{}:{} getting".format(var, self._blobs[var]))
+        try:
+            return self._blobs[var]
+        except:
+            print("WARNING: CANNOT FOUND blob {}".format(var))
+            return None
+
+
+log = TransLog()
+
+layer_names = {}
+
+
+def _conv2d(raw, input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    x = raw(input, weight, bias, stride, padding, dilation, groups)
+    name = log.add_layer(name='conv')
+    log.add_blobs([x], name='conv_blob')
+    layer = caffe_net.Layer_param(name=name, type='Convolution',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+    layer.conv_param(x.size()[1], weight.size()[2:], stride=_pair(stride),
+                     pad=_pair(padding), dilation=_pair(dilation), bias_term=bias is not None, groups=groups)
+    if bias is not None:
+        layer.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
+        #print('conv2d weight, bias: ',weight.cpu().data.numpy(), bias.cpu().data.numpy())
+
+    else:
+        layer.param.convolution_param.bias_term = False
+        layer.add_data(weight.cpu().data.numpy())
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _conv_transpose2d(raw, input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
+    x = raw(input, weight, bias, stride, padding, output_padding, groups, dilation)
+    name = log.add_layer(name='conv_transpose')
+    log.add_blobs([x], name='conv_transpose_blob')
+    layer = caffe_net.Layer_param(name=name, type='Deconvolution',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+    layer.conv_param(x.size()[1], weight.size()[2:], stride=_pair(stride),
+                     pad=_pair(padding), dilation=_pair(dilation), bias_term=bias is not None)
+    if bias is not None:
+        layer.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
+    else:
+        layer.param.convolution_param.bias_term = False
+        layer.add_data(weight.cpu().data.numpy())
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _linear(raw, input, weight, bias=None):
+    x = raw(input, weight, bias)
+    layer_name = log.add_layer(name='fc')
+    top_blobs = log.add_blobs([x], name='fc_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='InnerProduct',
+                                  bottom=[log.blobs(input)], top=top_blobs)
+    layer.fc_param(x.size()[1], has_bias=bias is not None)
+    if bias is not None:
+        layer.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
+    else:
+        layer.add_data(weight.cpu().data.numpy())
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _split(raw, tensor, split_size, dim=0):
+    # split in pytorch is slice in caffe
+    x = raw(tensor, split_size, dim)
+    layer_name = log.add_layer('split')
+    top_blobs = log.add_blobs(x, name='split_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Slice',
+                                  bottom=[log.blobs(tensor)], top=top_blobs)
+    slice_num = int(np.floor(tensor.size()[dim] / split_size))
+    slice_param = caffe_net.pb.SliceParameter(axis=dim, slice_point=[split_size * i for i in range(1, slice_num)])
+    layer.param.slice_param.CopyFrom(slice_param)
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _pool(type, raw, input, x, kernel_size, stride, padding, ceil_mode):
+    # TODO dilation,ceil_mode,return indices
+    layer_name = log.add_layer(name='{}_pool'.format(type))
+    top_blobs = log.add_blobs([x], name='{}_pool_blob'.format(type))
+    layer = caffe_net.Layer_param(name=layer_name, type='Pooling', bottom=[log.blobs(input)], top=top_blobs)
+
+    # TODO w,h different kernel, stride and padding
+    # processing ceil mode
+    layer.pool_param(kernel_size=kernel_size, stride=kernel_size if stride is None else stride,
+                     pad=padding, type=type.upper())
+    log.cnet.add_layer(layer)
+    if ceil_mode == False and stride is not None:
+        oheight = (input.size()[2] - _pair(kernel_size)[0] + 2 * _pair(padding)[0]) % (_pair(stride)[0])
+        owidth = (input.size()[3] - _pair(kernel_size)[1] + 2 * _pair(padding)[1]) % (_pair(stride)[1])
+        if oheight != 0 or owidth != 0:
+            caffe_out = raw(input, kernel_size, stride, padding, ceil_mode=False)
+            print("WARNING: the output shape miss match at {}: "
+
+                  "input {} output---Pytorch:{}---Caffe:{}\n"
+                  "This is caused by the different implementation that ceil mode in caffe and the floor mode in pytorch.\n"
+                  "You can add the clip layer in caffe prototxt manually if shape mismatch error is caused in caffe. ".format(
+                layer_name, input.size(), x.size(), caffe_out.size()))
+
+
+def _max_pool2d(raw, input, kernel_size, stride=None, padding=0, dilation=1,
+                ceil_mode=False, return_indices=False):
+    x = raw(input, kernel_size, stride, padding, dilation, ceil_mode, return_indices)
+    _pool('max', raw, input, x, kernel_size, stride, padding, ceil_mode)
+    return x
+
+
+def _avg_pool2d(raw, input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True):
+    x = raw(input, kernel_size, stride, padding, ceil_mode, count_include_pad)
+    _pool('ave', raw, input, x, kernel_size, stride, padding, ceil_mode)
+    return x
+
+
+def _max(raw, *args):
+    x = raw(*args)
+    if len(args) == 1:
+        # TODO max in one tensor
+        assert NotImplementedError
+    else:
+        bottom_blobs = []
+        for arg in args:
+            bottom_blobs.append(log.blobs(arg))
+        layer_name = log.add_layer(name='max')
+        top_blobs = log.add_blobs([x], name='max_blob')
+        layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                      bottom=bottom_blobs, top=top_blobs)
+        layer.param.eltwise_param.operation = 2
+        log.cnet.add_layer(layer)
+    return x
+
+
+def _cat(raw, inputs, dimension=0):
+    x = raw(inputs, dimension)
+    bottom_blobs = []
+    for input in inputs:
+        bottom_blobs.append(log.blobs(input))
+    layer_name = log.add_layer(name='cat')
+    top_blobs = log.add_blobs([x], name='cat_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Concat',
+                                  bottom=bottom_blobs, top=top_blobs)
+    layer.param.concat_param.axis = dimension
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _dropout(raw, input, p=0.5, training=False, inplace=False):
+    x = raw(input, p, training, inplace)
+    bottom_blobs = [log.blobs(input)]
+    layer_name = log.add_layer(name='dropout')
+    top_blobs = log.add_blobs([x], name=bottom_blobs[0], with_num=False)
+    layer = caffe_net.Layer_param(name=layer_name, type='Dropout',
+                                  bottom=bottom_blobs, top=top_blobs)
+    layer.param.dropout_param.dropout_ratio = p
+    layer.param.include.extend([caffe_net.pb.NetStateRule(phase=0)])  # 1 for test, 0 for train
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _threshold(raw, input, threshold, value, inplace=False):
+    # for threshold or relu
+    if threshold == 0 and value == 0:
+        x = raw(input, threshold, value, inplace)
+        bottom_blobs = [log.blobs(input)]
+        name = log.add_layer(name='relu')
+        log.add_blobs([x], name='relu_blob')
+        layer = caffe_net.Layer_param(name=name, type='ReLU',
+                                      bottom=bottom_blobs, top=[log.blobs(x)])
+        log.cnet.add_layer(layer)
+        return x
+    if value != 0:
+        raise NotImplemented("value !=0 not implemented in caffe")
+    x = raw(input, input, threshold, value, inplace)
+    bottom_blobs = [log.blobs(input)]
+    layer_name = log.add_layer(name='threshold')
+    top_blobs = log.add_blobs([x], name='threshold_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Threshold',
+                                  bottom=bottom_blobs, top=top_blobs)
+    layer.param.threshold_param.threshold = threshold
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _relu(raw, input, inplace=False):
+    # for threshold or prelu
+    x = raw(input, False)
+    name = log.add_layer(name='relu')
+    log.add_blobs([x], name='relu_blob')
+    layer = caffe_net.Layer_param(name=name, type='ReLU',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _prelu(raw, input, weight):
+    # for threshold or prelu
+    x = raw(input, weight)
+    bottom_blobs = [log.blobs(input)]
+    name = log.add_layer(name='prelu')
+    log.add_blobs([x], name='prelu_blob')
+    layer = caffe_net.Layer_param(name=name, type='PReLU',
+                                  bottom=bottom_blobs, top=[log.blobs(x)])
+    if weight.size()[0] == 1:
+        layer.param.prelu_param.channel_shared = True
+        layer.add_data(weight.cpu().data.numpy()[0])
+    else:
+        layer.add_data(weight.cpu().data.numpy())
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _leaky_relu(raw, input, negative_slope=0.01, inplace=False):
+    x = raw(input, negative_slope)
+    name = log.add_layer(name='leaky_relu')
+    log.add_blobs([x], name='leaky_relu_blob')
+    layer = caffe_net.Layer_param(name=name, type='ReLU',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+    layer.param.relu_param.negative_slope = negative_slope
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _tanh(raw, input):
+    # for tanh activation
+    x = raw(input)
+    name = log.add_layer(name='tanh')
+    log.add_blobs([x], name='tanh_blob')
+    layer = caffe_net.Layer_param(name=name, type='TanH',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _softmax(raw, input, dim=None, _stacklevel=3):
+    # for F.softmax
+    x = raw(input, dim=dim)
+    if dim is None:
+        dim = F._get_softmax_dim('softmax', input.dim(), _stacklevel)
+    bottom_blobs = [log.blobs(input)]
+    name = log.add_layer(name='softmax')
+    log.add_blobs([x], name='softmax_blob')
+    layer = caffe_net.Layer_param(name=name, type='Softmax',
+                                  bottom=bottom_blobs, top=[log.blobs(x)])
+    layer.param.softmax_param.axis = dim
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _sigmoid(raw, input):
+    # for tanh activation
+    x = raw(input)
+    name = log.add_layer(name='Sigmoid')
+    log.add_blobs([x], name='Sigmoid_blob')
+    layer = caffe_net.Layer_param(name=name, type='Sigmoid',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _batch_norm(raw, input, running_mean, running_var, weight=None, bias=None,
+                training=False, momentum=0.1, eps=1e-5):
+    # because the runing_mean and runing_var will be changed after the _batch_norm operation, we first save the parameters
+
+    x = raw(input, running_mean, running_var, weight, bias,
+            training, momentum, eps)
+    bottom_blobs = [log.blobs(input)]
+    layer_name1 = log.add_layer(name='batch_norm')
+    top_blobs = log.add_blobs([x], name='batch_norm_blob')
+    layer1 = caffe_net.Layer_param(name=layer_name1, type='BatchNorm',
+                                   bottom=bottom_blobs, top=top_blobs)
+    if running_mean is None or running_var is None:
+        # not use global_stats, normalization is performed over the current mini-batch
+        layer1.batch_norm_param(use_global_stats=0, eps=eps)
+    else:
+        layer1.batch_norm_param(use_global_stats=1, eps=eps)
+        running_mean_clone = running_mean.clone()
+        running_var_clone = running_var.clone()
+        layer1.add_data(running_mean_clone.cpu().numpy(), running_var_clone.cpu().numpy(), np.array([1.0]))
+        #print('running_mean: ',running_mean_clone.cpu().numpy())
+        #print('running_var: ',running_var_clone.cpu().numpy())
+    log.cnet.add_layer(layer1)
+    if weight is not None and bias is not None:
+        layer_name2 = log.add_layer(name='bn_scale')
+        layer2 = caffe_net.Layer_param(name=layer_name2, type='Scale',
+                                       bottom=top_blobs, top=top_blobs)
+        layer2.param.scale_param.bias_term = True
+        layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
+        log.cnet.add_layer(layer2)
+        #print('scale weight: ', weight.cpu().data.numpy())
+        #print('scale bias: ', bias.cpu().data.numpy())
+    return x
+
+
+def _instance_norm(raw, input, running_mean=None, running_var=None, weight=None,
+                   bias=None, use_input_stats=True, momentum=0.1, eps=1e-5):
+    # TODO: the batch size!=1 view operations
+    print("WARNING: The Instance Normalization transfers to Caffe using BatchNorm, so the batch size should be 1")
+    if running_var is not None or weight is not None:
+        # TODO: the affine=True or track_running_stats=True case
+        raise NotImplementedError("not implement the affine=True or track_running_stats=True case InstanceNorm")
+    x = torch.batch_norm(
+        input, weight, bias, running_mean, running_var,
+        use_input_stats, momentum, eps, torch.backends.cudnn.enabled)
+    bottom_blobs = [log.blobs(input)]
+    layer_name1 = log.add_layer(name='instance_norm')
+    top_blobs = log.add_blobs([x], name='instance_norm_blob')
+    layer1 = caffe_net.Layer_param(name=layer_name1, type='BatchNorm',
+                                   bottom=bottom_blobs, top=top_blobs)
+    if running_mean is None or running_var is None:
+        # not use global_stats, normalization is performed over the current mini-batch
+        layer1.batch_norm_param(use_global_stats=0, eps=eps)
+        running_mean = torch.zeros(input.size()[1])
+        running_var = torch.ones(input.size()[1])
+    else:
+        layer1.batch_norm_param(use_global_stats=1, eps=eps)
+    running_mean_clone = running_mean.clone()
+    running_var_clone = running_var.clone()
+    layer1.add_data(running_mean_clone.cpu().numpy(), running_var_clone.cpu().numpy(), np.array([1.0]))
+    log.cnet.add_layer(layer1)
+    if weight is not None and bias is not None:
+        layer_name2 = log.add_layer(name='bn_scale')
+        layer2 = caffe_net.Layer_param(name=layer_name2, type='Scale',
+                                       bottom=top_blobs, top=top_blobs)
+        layer2.param.scale_param.bias_term = True
+        layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
+        log.cnet.add_layer(layer2)
+    return x
+
+
+# upsample layer
+def _interpolate(raw, input, size=None, scale_factor=None, mode='nearest', align_corners=None):
+    # 定义的参数包括 scale,即输出与输入的尺寸比例,如 2;scale_h、scale_w,
+    # 同 scale,分别为 h、w 方向上的尺寸比例;pad_out_h、pad_out_w,仅在 scale 为 2 时
+    # 有用,对输出进行额外 padding 在 h、w 方向上的数值;upsample_h、upsample_w,输
+    # 出图像尺寸的数值。在 Upsample 的相关代码中,推荐仅仅使用 upsample_h、
+    # upsample_w 准确定义 Upsample 层的输出尺寸,其他所有的参数都不推荐继续使用。
+    '''
+    if mode == 'bilinear':      
+        x = raw(input, size, scale_factor, mode)
+        name = log.add_layer(name='conv_transpose')
+        log.add_blobs([x], name='conv_transpose_blob')
+        layer = caffe_net.Layer_param(name=name, type='Deconvolution',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+        print('Deconv: ', name)
+        print(input.shape)
+        print(x.size())
+        print(size)
+        factor = float(size[0]) / input.shape[2]
+        C = x.size()[1]
+        print(factor,C) 
+        kernel_size = int(2 * factor - factor % 2)
+        stride = int(factor)
+        num_output = C
+        group = C
+        pad = math.ceil((factor-1) / 2.)
+        print('kernel_size, stride, num_output, group, pad')
+        print(kernel_size, stride, num_output, group, pad)
+        layer.conv_param(num_output, kernel_size, stride=stride,
+                     pad=pad, weight_filler_type='bilinear', bias_term=False, groups=group)
+
+        layer.param.convolution_param.bias_term = False
+        log.cnet.add_layer(layer)
+        return x
+    '''
+    # transfer bilinear align_corners=True to caffe-interp
+    if mode == "bilinear" and align_corners == True:
+        x = raw(input, size, scale_factor, mode)
+        name = log.add_layer(name='interp')
+        log.add_blobs([x], name='interp_blob')
+        layer = caffe_net.Layer_param(name=name, type='Interp',
+                                  bottom=[log.blobs(input)], top=[log.blobs(x)])
+        layer.interp_param(size=size, scale_factor=scale_factor)
+        log.cnet.add_layer(layer)
+        return x
+
+    # for nearest _interpolate
+    if mode != "nearest" or align_corners != None:
+        raise NotImplementedError("not implement F.interpolate totoaly")
+    x = raw(input, size, scale_factor, mode)
+    layer_name = log.add_layer(name='upsample')
+    top_blobs = log.add_blobs([x], name='upsample_blob'.format(type))
+    layer = caffe_net.Layer_param(name=layer_name, type='Upsample',
+                                  bottom=[log.blobs(input)], top=top_blobs)
+    #layer.upsample_param(size=(input.size(2), input.size(3)), scale_factor=scale_factor)
+    #layer.upsample_param(size=size, scale_factor=scale_factor)
+    layer.upsample_param(size=None, scale_factor=size[0])
+    
+    log.cnet.add_layer(layer)
+    return x
+
+
+# ----- for Variable operations --------
+
+def _view(input, *args):
+    x = raw_view(input, *args)
+    if not NET_INITTED:
+        return x
+    layer_name = log.add_layer(name='view')
+    top_blobs = log.add_blobs([x], name='view_blob')
+
+    # print('*'*60)
+    # print('input={}'.format(input))
+    # print('layer_name={}'.format(layer_name))
+    # print('top_blobs={}'.format(top_blobs))
+
+    layer = caffe_net.Layer_param(name=layer_name, type='Reshape', bottom=[log.blobs(input)], top=top_blobs)
+    # TODO: reshpae added to nn_tools layer
+    dims = list(args)
+    dims[0] = 0  # the first dim should be batch_size
+    layer.param.reshape_param.shape.CopyFrom(caffe_net.pb.BlobShape(dim=dims))
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _mean(input, *args, **kwargs):
+    x = raw_mean(input, *args, **kwargs)
+    if not NET_INITTED:
+        return x
+    layer_name = log.add_layer(name='mean')
+    top_blobs = log.add_blobs([x], name='mean_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Reduction',
+                                  bottom=[log.blobs(input)], top=top_blobs)
+    if len(args) == 1:
+        dim = args[0]
+    elif 'dim' in kwargs:
+        dim = kwargs['dim']
+    else:
+        raise NotImplementedError('mean operation must specify a dim')
+    layer.param.reduction_param.operation = 4
+    layer.param.reduction_param.axis = dim
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _add(input, *args):
+    # check if add a const value
+    if isinstance(args[0], int):
+        print('value: ',args[0])
+        x = raw__add__(input, *args)
+        #x = raw(input)
+        layer_name = log.add_layer(name='scale')
+        log.add_blobs([x], name='Scale_blob')
+        layer = caffe_net.Layer_param(name=layer_name, type='Scale',
+                                       bottom=[log.blobs(input)], top=[log.blobs(x)])
+        dim = x.shape[1]
+        layer.param.scale_param.bias_term = True
+        weight = np.ones(dim, dtype=np.float32)
+        bias = args[0] * np.ones(dim, dtype=np.float32)
+        layer.add_data(weight, bias)
+        log.cnet.add_layer(layer)
+        return x
+    # otherwise add a tensor
+    x = raw__add__(input, *args)
+    if not NET_INITTED:
+        return x
+    layer_name = log.add_layer(name='add')
+    top_blobs = log.add_blobs([x], name='add_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                  bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs)
+    layer.param.eltwise_param.operation = 1  # sum is 1
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _iadd(input, *args):
+    x = raw__iadd__(input, *args)
+    if not NET_INITTED:
+        return x
+    x = x.clone()
+    layer_name = log.add_layer(name='add')
+    top_blobs = log.add_blobs([x], name='add_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                  bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs)
+    layer.param.eltwise_param.operation = 1  # sum is 1
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _sub(input, *args):
+    x = raw__sub__(input, *args)
+    if not NET_INITTED:
+        return x
+    layer_name = log.add_layer(name='sub')
+    top_blobs = log.add_blobs([x], name='sub_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                  bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs)
+    layer.param.eltwise_param.operation = 1  # sum is 1
+    layer.param.eltwise_param.coeff.extend([1., -1.])
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _isub(input, *args):
+    x = raw__isub__(input, *args)
+    if not NET_INITTED:
+        return x
+    x = x.clone()
+    layer_name = log.add_layer(name='sub')
+    top_blobs = log.add_blobs([x], name='sub_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                  bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs)
+    layer.param.eltwise_param.operation = 1  # sum is 1
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _mul(input, *args):
+    x = raw__sub__(input, *args)
+    if not NET_INITTED:
+        return x
+    layer_name = log.add_layer(name='mul')
+    top_blobs = log.add_blobs([x], name='mul_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                  bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs)
+    layer.param.eltwise_param.operation = 0  # product is 1
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _imul(input, *args):
+    x = raw__isub__(input, *args)
+    if not NET_INITTED:
+        return x
+    x = x.clone()
+    layer_name = log.add_layer(name='mul')
+    top_blobs = log.add_blobs([x], name='mul_blob')
+    layer = caffe_net.Layer_param(name=layer_name, type='Eltwise',
+                                  bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs)
+    layer.param.eltwise_param.operation = 0  # product is 1
+    layer.param.eltwise_param.coeff.extend([1., -1.])
+    log.cnet.add_layer(layer)
+    return x
+
+
+def _adaptive_avg_pool2d(raw, input, output_size):
+    _output_size = _list_with_default(output_size, input.size())
+    x = raw(input, _output_size)
+    _pool('ave', raw, input, x, input.shape[2], input.shape[2], 0, False)
+    return x
+
+
+# 核心组件，通过该类，实现对torch的function中的operators的输入，输出以及参数的读取
+class Rp(object):
+    def __init__(self, raw, replace, **kwargs):
+        # replace the raw function to replace function
+        self.obj = replace
+        self.raw = raw
+
+    def __call__(self, *args, **kwargs):
+        if not NET_INITTED:
+            return self.raw(*args, **kwargs)
+        for stack in traceback.walk_stack(None):
+            if 'self' in stack[0].f_locals:
+                layer = stack[0].f_locals['self']
+                if layer in layer_names:
+                    log.pytorch_layer_name = layer_names[layer]
+                    print(layer_names[layer])
+                    break
+        out = self.obj(self.raw, *args, **kwargs)
+        # if isinstance(out,Variable):
+        #     out=[out]
+        return out
+
+
+F.conv2d = Rp(F.conv2d, _conv2d)
+F.linear = Rp(F.linear, _linear)
+F.relu = Rp(F.relu, _relu)
+
+F.leaky_relu = Rp(F.leaky_relu, _leaky_relu)
+F.max_pool2d = Rp(F.max_pool2d, _max_pool2d)
+F.avg_pool2d = Rp(F.avg_pool2d, _avg_pool2d)
+F.dropout = Rp(F.dropout, _dropout)
+F.threshold = Rp(F.threshold, _threshold)
+F.prelu = Rp(F.prelu, _prelu)
+F.batch_norm = Rp(F.batch_norm, _batch_norm)
+F.instance_norm = Rp(F.instance_norm, _instance_norm)
+F.softmax = Rp(F.softmax, _softmax)
+F.conv_transpose2d = Rp(F.conv_transpose2d, _conv_transpose2d)
+F.interpolate = Rp(F.interpolate, _interpolate)
+F.adaptive_avg_pool2d = Rp(F.adaptive_avg_pool2d, _adaptive_avg_pool2d)
+
+torch.split = Rp(torch.split, _split)
+torch.max = Rp(torch.max, _max)
+torch.cat = Rp(torch.cat, _cat)
+torch.sigmoid = Rp(torch.sigmoid, _sigmoid)
+
+# TODO: other types of the view function
+try:
+    raw_view = Variable.view
+    Variable.view = _view
+    raw_mean = Variable.mean
+    Variable.mean = _mean
+    raw__add__ = Variable.__add__
+    Variable.__add__ = _add
+    raw__iadd__ = Variable.__iadd__
+    Variable.__iadd__ = _iadd
+    raw__sub__ = Variable.__sub__
+    Variable.__sub__ = _sub
+    raw__isub__ = Variable.__isub__
+    Variable.__isub__ = _isub
+    raw__mul__ = Variable.__mul__
+    Variable.__mul__ = _mul
+    raw__imul__ = Variable.__imul__
+    Variable.__imul__ = _imul
+except:
+    # for new version 0.4.0 and later version
+    for t in [torch.Tensor]:
+        raw_view = t.view
+        t.view = _view
+        raw_mean = t.mean
+        t.mean = _mean
+        raw__add__ = t.__add__
+        t.__add__ = _add
+        raw__iadd__ = t.__iadd__
+        t.__iadd__ = _iadd
+        raw__sub__ = t.__sub__
+        t.__sub__ = _sub
+        raw__isub__ = t.__isub__
+        t.__isub__ = _isub
+        raw__mul__ = t.__mul__
+        t.__mul__ = _mul
+        raw__imul__ = t.__imul__
+        t.__imul__ = _imul
+
+
+def trans_net(net, input_var, name='TransferedPytorchModel'):
+    print('Starting Transform, This will take a while')
+    log.init([input_var])
+    log.cnet.net.name = name
+    log.cnet.net.input.extend([log.blobs(input_var)])
+    log.cnet.net.input_dim.extend(input_var.size())
+    global NET_INITTED
+    NET_INITTED = True
+    for name, layer in net.named_modules():
+        layer_names[layer] = name
+    print("torch ops name:", layer_names)
+    out = net.forward(input_var)
+    print('Transform Completed')
+
+
+def save_prototxt(save_name):
+    log.cnet.save_prototxt(save_name)
+
+
+def save_caffemodel(save_name):
+    log.cnet.save(save_name)

CION_ReIDZoo/CION_Finetune/FastReID/tools/deploy/trt_export.py

@@ -0,0 +1,82 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: [email protected]
+"""
+
+import argparse
+import os
+import numpy as np
+import sys
+
+sys.path.append('../../')
+sys.path.append("/export/home/lxy/runtimelib-tensorrt-tiny/build")
+
+import pytrt
+from fastreid.utils.logger import setup_logger
+from fastreid.utils.file_io import PathManager
+
+
+logger = setup_logger(name='trt_export')
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Convert ONNX to TRT model")
+
+    parser.add_argument(
+        "--name",
+        default="baseline",
+        help="name for converted model"
+    )
+    parser.add_argument(
+        "--output",
+        default='outputs/trt_model',
+        help='path to save converted trt model'
+    )
+    parser.add_argument(
+        "--onnx-model",
+        default='outputs/onnx_model/baseline.onnx',
+        help='path to onnx model'
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=256,
+        help="height of image"
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=128,
+        help="width of image"
+    )
+    return parser
+
+
+def export_trt_model(onnxModel, engineFile, input_numpy_array):
+    r"""
+    Export a model to trt format.
+    """
+
+    trt = pytrt.Trt()
+
+    customOutput = []
+    maxBatchSize = 1
+    calibratorData = []
+    mode = 2
+    trt.CreateEngine(onnxModel, engineFile, customOutput, maxBatchSize, mode, calibratorData)
+    trt.DoInference(input_numpy_array)  # slightly different from c++
+    return 0
+
+
+if __name__ == '__main__':
+    args = get_parser().parse_args()
+
+    inputs = np.zeros(shape=(32, args.height, args.width, 3))
+    onnxModel = args.onnx_model
+    engineFile = os.path.join(args.output, args.name+'.engine')
+
+    PathManager.mkdirs(args.output)
+    export_trt_model(onnxModel, engineFile, inputs)
+
+    logger.info(f"Export trt model in {args.output} successfully!")

open_clip/.github/workflows/ci.yml

@@ -0,0 +1,121 @@
+name: Continuous integration
+
+on:
+  push:
+    branches:
+    - main 
+    paths-ignore:
+    - '**.md'
+    - 'CITATION.cff'
+    - 'LICENSE'
+    - '.gitignore'
+    - 'docs/**'
+  pull_request:
+    branches:
+    - main 
+    paths-ignore:
+    - '**.md'
+    - 'CITATION.cff'
+    - 'LICENSE'
+    - '.gitignore'
+    - 'docs/**'
+  workflow_dispatch:
+    inputs:
+      manual_revision_reference:
+        required: false
+        type: string
+      manual_revision_test:
+        required: false
+        type: string
+
+env:
+  REVISION_REFERENCE: v2.8.2
+  #9d31b2ec4df6d8228f370ff20c8267ec6ba39383 earliest compatible v2.7.0 + pretrained_hf param
+
+jobs:
+  Tests:
+    strategy:
+      matrix:
+        os: [ ubuntu-latest ] #, macos-latest ]
+        python: [ 3.8 ]
+        job_num: [ 4 ]
+        job: [ 1, 2, 3, 4 ]
+    runs-on: ${{ matrix.os }}
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+        ref: ${{ inputs.manual_revision_test }}
+    - name: Set up Python ${{ matrix.python }}
+      id: pythonsetup
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python }}
+    - name: Venv cache
+      id: venv-cache
+      uses: actions/cache@v3
+      with:
+        path: .env
+        key: venv-${{ matrix.os }}-${{ steps.pythonsetup.outputs.python-version }}-${{ hashFiles('requirements*') }}
+    - name: Pytest durations cache
+      uses: actions/cache@v3
+      with:
+        path: .test_durations
+        key: test_durations-${{ matrix.os }}-${{ steps.pythonsetup.outputs.python-version }}-${{ matrix.job }}-${{ github.run_id }}
+        restore-keys: test_durations-0-
+    - name: Setup
+      if: steps.venv-cache.outputs.cache-hit != 'true'
+      run: |
+        python3 -m venv .env
+        source .env/bin/activate
+        pip install -e .[test]
+    - name: Prepare test data
+      run: |
+        source .env/bin/activate
+        python -m pytest \
+          --quiet --co \
+          --splitting-algorithm least_duration \
+          --splits ${{ matrix.job_num }} \
+          --group ${{ matrix.job }} \
+          -m regression_test \
+          tests \
+          | head -n -2 | grep -Po 'test_inference_with_data\[\K[^]]*(?=-False]|-True])' \
+          > models_gh_runner.txt
+        if [ -n "${{ inputs.manual_revision_reference }}" ]; then
+          REVISION_REFERENCE=${{ inputs.manual_revision_reference }}
+        fi
+        python tests/util_test.py \
+          --save_model_list models_gh_runner.txt \
+          --model_list models_gh_runner.txt \
+          --git_revision $REVISION_REFERENCE
+    - name: Unit tests
+      run: |
+        source .env/bin/activate
+        if [[ -f .test_durations ]]
+        then
+          cp .test_durations durations_1
+          mv .test_durations durations_2
+        fi
+        python -m pytest \
+          -x -s -v \
+          --splitting-algorithm least_duration \
+          --splits ${{ matrix.job_num }} \
+          --group ${{ matrix.job }} \
+          --store-durations \
+          --durations-path durations_1 \
+          --clean-durations \
+          -m "not regression_test" \
+          tests
+        OPEN_CLIP_TEST_REG_MODELS=models_gh_runner.txt python -m pytest \
+          -x -s -v \
+          --store-durations \
+          --durations-path durations_2 \
+          --clean-durations \
+          -m "regression_test" \
+          tests
+        jq -s -S 'add' durations_* > .test_durations
+    - name: Collect pytest durations
+      uses: actions/upload-artifact@v3
+      with:
+        name: pytest_durations_${{ matrix.os }}-${{ matrix.python }}-${{ matrix.job }}
+        path: .test_durations

open_clip/.github/workflows/clear-cache.yml

@@ -0,0 +1,29 @@
+name: Clear cache
+
+on:
+  workflow_dispatch:
+
+permissions:
+  actions: write
+
+jobs:
+  clear-cache:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Clear cache
+        uses: actions/github-script@v6
+        with:
+          script: |
+            const caches = await github.rest.actions.getActionsCacheList({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+            })
+            for (const cache of caches.data.actions_caches) {
+              console.log(cache)
+              await github.rest.actions.deleteActionsCacheById({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                cache_id: cache.id,
+              })
+            }
+

open_clip/.github/workflows/python-publish.yml

@@ -0,0 +1,37 @@
+name: Release
+
+on:
+  push:
+    branches:
+    - main
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions-ecosystem/action-regex-match@v2
+      id: regex-match
+      with:
+        text: ${{ github.event.head_commit.message }}
+        regex: '^Release ([^ ]+)'
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.8'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine build
+    - name: Release
+      if: ${{ steps.regex-match.outputs.match != '' }}
+      uses: softprops/action-gh-release@v1
+      with:
+        tag_name: v${{ steps.regex-match.outputs.group1 }}
+    - name: Build and publish
+      if: ${{ steps.regex-match.outputs.match != '' }}
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+      run: |
+        python -m build
+        twine upload dist/*

open_clip/docs/openclip_retrieval_results.csv

@@ -0,0 +1,122 @@
+name,pretrained,params (M),FLOPs (B),Average score,Flickr image retr. R@1,Flickr image retr. R@5,Flickr image retr. R@10,Flickr text retr. R@1,Flickr text retr. R@5,Flickr text retr. R@10,MSCOCO image retr. R@1,MSCOCO image retr. R@5,MSCOCO image retr. R@10,MSCOCO text retr. R@1,MSCOCO text retr. R@5,MSCOCO text retr. R@10,WinoGAViL avg jaccard score,WinoGAViL jaccard score 10,WinoGAViL jaccard score 10-12,WinoGAViL jaccard score 12,WinoGAViL jaccard score 5,WinoGAViL jaccard score 5-6,WinoGAViL jaccard score 6
+ViT-SO400M-14-SigLIP-384,webli,877.96,723.48,0.7721,0.8296,0.9610,0.9804,0.9430,0.9970,0.9980,0.5421,0.7678,0.8424,0.7242,0.8998,0.9448,0.6181,0.5807,0.5754,0.5701,0.6427,0.6316,0.6210
+ViT-H-14-378-quickgelu,dfn5b,986.71,1054.05,0.7719,0.8202,0.9598,0.9798,0.9400,0.9920,0.9960,0.5564,0.7920,0.8626,0.7188,0.9048,0.9496,0.6123,0.5668,0.5674,0.5679,0.6409,0.6265,0.6127
+ViT-L-16-SigLIP-384,webli,652.48,422.91,0.7642,0.8142,0.9536,0.9748,0.9370,0.9920,0.9990,0.5391,0.7657,0.8399,0.7190,0.9006,0.9410,0.6070,0.5562,0.5550,0.5539,0.6422,0.6234,0.6056
+ViT-H-14-quickgelu,dfn5b,986.11,381.68,0.7640,0.8012,0.9526,0.9738,0.9280,0.9940,0.9990,0.5391,0.7805,0.8553,0.7230,0.9024,0.9452,0.6022,0.5568,0.5561,0.5554,0.6315,0.6167,0.6028
+ViT-L-16-SigLIP-256,webli,652.15,201.62,0.7619,0.7904,0.9446,0.9684,0.9180,0.9900,0.9980,0.5228,0.7580,0.8334,0.7080,0.8870,0.9374,0.6142,0.5877,0.5748,0.5619,0.6450,0.6267,0.6092
+ViT-SO400M-14-SigLIP,webli,877.36,233.54,0.7567,0.7526,0.9226,0.9554,0.9100,0.9910,0.9980,0.5176,0.7527,0.8300,0.6966,0.8908,0.9348,0.6189,0.5736,0.5665,0.5594,0.6430,0.6354,0.6282
+ViT-B-16-SigLIP-512,webli,203.79,227.26,0.7554,0.7906,0.9458,0.9690,0.9250,0.9920,0.9960,0.5055,0.7421,0.8217,0.6872,0.8786,0.9266,0.6070,0.5761,0.5696,0.5633,0.6291,0.6187,0.6088
+ViT-B-16-SigLIP-384,webli,203.45,123.15,0.7542,0.7848,0.9416,0.9708,0.9270,0.9910,0.9980,0.4990,0.7366,0.8204,0.6774,0.8758,0.9258,0.6077,0.5730,0.5719,0.5709,0.6314,0.6190,0.6072
+ViT-B-16-SigLIP-256,webli,203.2,57.84,0.7462,0.7504,0.9242,0.9626,0.9040,0.9830,0.9930,0.4834,0.7254,0.8099,0.6614,0.8634,0.9182,0.6113,0.5754,0.5715,0.5677,0.6373,0.6238,0.6110
+ViT-B-16-SigLIP,webli,203.16,46.44,0.7442,0.7468,0.9230,0.9562,0.8910,0.9800,0.9930,0.4778,0.7244,0.8100,0.6574,0.8542,0.9126,0.6134,0.5726,0.5738,0.5750,0.6335,0.6259,0.6186
+coca_ViT-L-14,mscoco_finetuned_laion2b_s13b_b90k,638.45,214.52,0.7432,0.7846,0.9452,0.9712,0.8900,0.9870,0.9950,0.5374,0.7779,0.8567,0.6682,0.8762,0.9242,0.5762,0.5265,0.5146,0.5028,0.5978,0.5956,0.5934
+ViT-H-14-quickgelu,metaclip_fullcc,986.11,381.68,0.7412,0.7834,0.9464,0.9692,0.9180,0.9870,0.9970,0.4882,0.7323,0.8136,0.6622,0.8618,0.9188,0.5891,0.5316,0.5312,0.5308,0.6220,0.6073,0.5934
+ViT-L-14-quickgelu,dfn2b,427.62,175.33,0.7406,0.7546,0.9280,0.9608,0.8960,0.9850,0.9930,0.4856,0.7381,0.8237,0.6560,0.8584,0.9120,0.5955,0.5688,0.5519,0.5351,0.6219,0.6093,0.5973
+xlm-roberta-large-ViT-H-14,frozen_laion5b_s13b_b90k,1193.01,671.01,0.7363,0.7742,0.9392,0.9670,0.9180,0.9930,0.9980,0.4921,0.7305,0.8159,0.6596,0.8618,0.9184,0.5767,0.5277,0.5206,0.5135,0.6076,0.5943,0.5818
+ViT-B-16-SigLIP-i18n-256,webli,370.63,57.84,0.7347,0.7216,0.9036,0.9470,0.8960,0.9820,0.9910,0.4492,0.6948,0.7837,0.6448,0.8434,0.9060,0.6113,0.5887,0.5710,0.5534,0.6303,0.6240,0.6180
+ViT-L-14-quickgelu,metaclip_fullcc,427.62,175.33,0.7325,0.7642,0.9366,0.9646,0.9010,0.9850,0.9930,0.4709,0.7141,0.8026,0.6442,0.8504,0.9130,0.5817,0.5539,0.5357,0.5176,0.5990,0.5963,0.5937
+ViT-H-14-CLIPA-336,laion2b,968.64,800.88,0.7289,0.7814,0.9392,0.9684,0.9230,0.9920,0.9970,0.5013,0.7387,0.8224,0.6752,0.8802,0.9308,0.5452,0.4970,0.4878,0.4786,0.5764,0.5633,0.5508
+convnext_xxlarge,laion2b_s34b_b82k_augreg_rewind,1200.58,443.03,0.7238,0.7886,0.9438,0.9714,0.9110,0.9910,0.9970,0.5003,0.7407,0.8240,0.6740,0.8738,0.9264,0.5327,0.4826,0.4741,0.4656,0.5674,0.5512,0.5357
+EVA02-L-14-336,merged2b_s6b_b61k,428.08,395.16,0.7235,0.7802,0.9432,0.9684,0.8960,0.9890,0.9960,0.4794,0.7172,0.7999,0.6416,0.8522,0.9084,0.5532,0.5081,0.5053,0.5025,0.5719,0.5683,0.5649
+convnext_xxlarge,laion2b_s34b_b82k_augreg_soup,1200.58,443.03,0.7232,0.7902,0.9460,0.9706,0.9140,0.9920,0.9970,0.4997,0.7403,0.8227,0.6736,0.8722,0.9260,0.5316,0.4790,0.4702,0.4615,0.5651,0.5509,0.5374
+EVA01-g-14-plus,merged2b_s11b_b114k,1366.62,581.15,0.7231,0.7900,0.9448,0.9692,0.9170,0.9930,0.9970,0.5033,0.7398,0.8206,0.6816,0.8750,0.9250,0.5298,0.4787,0.4684,0.4582,0.5655,0.5491,0.5336
+ViT-bigG-14-CLIPA-336,datacomp1b,2517.76,2271.58,0.7220,0.7828,0.9408,0.9674,0.9220,0.9940,0.9980,0.5042,0.7432,0.8238,0.6776,0.8792,0.9306,0.5323,0.4624,0.4487,0.4350,0.5619,0.5587,0.5556
+ViT-bigG-14,laion2b_s39b_b160k,2539.57,1065.36,0.7213,0.7956,0.9498,0.9714,0.9290,0.9930,0.9980,0.5137,0.7490,0.8300,0.6738,0.8686,0.9266,0.5224,0.4714,0.4488,0.4263,0.5588,0.5456,0.5330
+EVA02-L-14,merged2b_s4b_b131k,427.76,175.3,0.7207,0.7732,0.9374,0.9678,0.8970,0.9850,0.9910,0.4747,0.7123,0.7971,0.6366,0.8434,0.9038,0.5518,0.5128,0.5081,0.5034,0.5767,0.5656,0.5551
+ViT-L-14,commonpool_xl_laion_s13b_b90k,427.62,175.33,0.7206,0.7490,0.9308,0.9630,0.8770,0.9900,0.9950,0.4677,0.7142,0.8042,0.6350,0.8458,0.9122,0.5635,0.5193,0.4966,0.4740,0.6023,0.5846,0.5677
+convnext_xxlarge,laion2b_s34b_b82k_augreg,1200.58,443.03,0.7201,0.7902,0.9414,0.9726,0.9160,0.9920,0.9970,0.4970,0.7400,0.8199,0.6698,0.8744,0.9236,0.5290,0.4730,0.4536,0.4343,0.5622,0.5527,0.5437
+ViT-H-14-CLIPA-336,datacomp1b,968.64,800.88,0.7192,0.7652,0.9350,0.9630,0.9060,0.9920,0.9950,0.4928,0.7309,0.8144,0.6716,0.8746,0.9276,0.5365,0.4801,0.4587,0.4375,0.5658,0.5610,0.5566
+EVA02-E-14-plus,laion2b_s9b_b144k,5044.89,2362.19,0.7188,0.7886,0.9434,0.9698,0.9410,0.9930,0.9980,0.5110,0.7492,0.8276,0.6872,0.8760,0.9274,0.5134,0.4576,0.4403,0.4232,0.5470,0.5364,0.5263
+ViT-H-14-CLIPA,datacomp1b,968.24,354.02,0.7178,0.7588,0.9330,0.9628,0.9100,0.9900,0.9960,0.4910,0.7291,0.8140,0.6698,0.8730,0.9272,0.5343,0.4793,0.4578,0.4365,0.5641,0.5584,0.5529
+ViT-bigG-14-CLIPA,datacomp1b,2517.22,1007.93,0.7175,0.7786,0.9374,0.9650,0.9190,0.9930,0.9980,0.4996,0.7414,0.8214,0.6780,0.8742,0.9312,0.5247,0.4461,0.4383,0.4306,0.5552,0.5519,0.5488
+ViT-B-16,dfn2b,149.62,41.09,0.7158,0.6912,0.8982,0.9406,0.8540,0.9780,0.9860,0.4339,0.6924,0.7882,0.6038,0.8310,0.8946,0.5852,0.5509,0.5417,0.5327,0.6058,0.5989,0.5923
+EVA02-E-14,laion2b_s4b_b115k,4704.59,2311.42,0.7152,0.7810,0.9438,0.9700,0.9220,0.9950,0.9970,0.5037,0.7429,0.8224,0.6746,0.8738,0.9256,0.5110,0.4316,0.4429,0.4542,0.5485,0.5325,0.5173
+convnext_large_d_320,laion2b_s29b_b131k_ft_soup,351.77,157.98,0.7148,0.7704,0.9394,0.9648,0.9230,0.9890,0.9980,0.4829,0.7251,0.8103,0.6502,0.8608,0.9180,0.5285,0.4574,0.4549,0.4525,0.5650,0.5517,0.5391
+ViT-L-14-quickgelu,metaclip_400m,427.62,175.33,0.7137,0.7342,0.9228,0.9574,0.8620,0.9800,0.9960,0.4381,0.6862,0.7780,0.6000,0.8290,0.8934,0.5701,0.5322,0.5175,0.5029,0.5941,0.5867,0.5796
+convnext_large_d_320,laion2b_s29b_b131k_ft,351.77,157.98,0.7122,0.7690,0.9364,0.9668,0.9140,0.9890,0.9980,0.4814,0.7232,0.8103,0.6384,0.8558,0.9168,0.5250,0.4727,0.4558,0.4389,0.5623,0.5468,0.5321
+convnext_large_d,laion2b_s26b_b102k_augreg,351.77,107.5,0.7102,0.7588,0.9310,0.9652,0.9180,0.9850,0.9940,0.4701,0.7139,0.8013,0.6400,0.8506,0.9140,0.5252,0.4835,0.4646,0.4457,0.5571,0.5443,0.5322
+ViT-L-14-CLIPA-336,datacomp1b,414.54,387.39,0.7092,0.7462,0.9240,0.9590,0.9040,0.9920,0.9980,0.4715,0.7159,0.8024,0.6564,0.8636,0.9180,0.5276,0.4601,0.4449,0.4297,0.5611,0.5537,0.5467
+ViT-L-14,datacomp_xl_s13b_b90k,427.62,175.33,0.7091,0.7338,0.9174,0.9554,0.8900,0.9860,0.9970,0.4573,0.7003,0.7916,0.6330,0.8414,0.9040,0.5471,0.4836,0.4666,0.4497,0.5912,0.5724,0.5546
+ViT-g-14,laion2b_s12b_b42k,1366.68,581.15,0.7085,0.7642,0.9364,0.9624,0.9090,0.9910,0.9980,0.4802,0.7238,0.8079,0.6492,0.8530,0.9152,0.5178,0.4498,0.4427,0.4357,0.5589,0.5414,0.5249
+coca_ViT-L-14,laion2b_s13b_b90k,638.45,214.52,0.7085,0.7428,0.9202,0.9542,0.8840,0.9930,0.9990,0.4565,0.7042,0.7921,0.6292,0.8370,0.9038,0.5403,0.4800,0.4739,0.4677,0.5704,0.5612,0.5525
+ViT-B-16-quickgelu,metaclip_fullcc,149.62,41.09,0.7077,0.7072,0.9082,0.9454,0.8550,0.9740,0.9890,0.4134,0.6718,0.7696,0.5936,0.8058,0.8784,0.5787,0.5391,0.5221,0.5052,0.6064,0.5965,0.5871
+ViT-H-14,laion2b_s32b_b79k,986.11,381.68,0.7057,0.7764,0.9418,0.9660,0.9070,0.9920,0.9970,0.4948,0.7338,0.8151,0.6592,0.8606,0.9188,0.4998,0.4289,0.4247,0.4206,0.5462,0.5234,0.5018
+convnext_base_w,laion_aesthetic_s13b_b82k,179.39,49.38,0.7047,0.7306,0.9160,0.9532,0.8880,0.9820,0.9930,0.4355,0.6858,0.7808,0.6120,0.8338,0.8978,0.5461,0.4937,0.4764,0.4592,0.5829,0.5681,0.5540
+ViT-L-14-CLIPA,datacomp1b,414.21,167.5,0.7044,0.7356,0.9206,0.9550,0.9020,0.9860,0.9930,0.4695,0.7084,0.7981,0.6512,0.8544,0.9150,0.5244,0.4495,0.4388,0.4282,0.5583,0.5514,0.5449
+ViT-g-14,laion2b_s34b_b88k,1366.68,581.15,0.7038,0.7772,0.9416,0.9690,0.9140,0.9930,0.9960,0.4878,0.7334,0.8151,0.6638,0.8598,0.9184,0.4973,0.4207,0.4104,0.4001,0.5458,0.5247,0.5047
+convnext_base_w_320,laion_aesthetic_s13b_b82k,179.39,71.94,0.7030,0.7402,0.9200,0.9524,0.9070,0.9830,0.9920,0.4387,0.6920,0.7824,0.6106,0.8312,0.8964,0.5385,0.4742,0.4610,0.4478,0.5695,0.5630,0.5567
+ViT-L-14,laion2b_s32b_b82k,427.62,175.33,0.7010,0.7552,0.9286,0.9594,0.8950,0.9870,0.9940,0.4651,0.7108,0.7980,0.6336,0.8398,0.9074,0.5142,0.4448,0.4385,0.4322,0.5507,0.5381,0.5261
+EVA01-g-14,laion400m_s11b_b41k,1136.44,547.36,0.7003,0.7264,0.9162,0.9510,0.8810,0.9830,0.9930,0.4406,0.6849,0.7732,0.6180,0.8328,0.8998,0.5363,0.4646,0.4640,0.4634,0.5665,0.5591,0.5521
+convnext_base_w,laion2b_s13b_b82k,179.39,49.38,0.6976,0.7214,0.9152,0.9538,0.8770,0.9800,0.9930,0.4285,0.6805,0.7770,0.6058,0.8324,0.8954,0.5311,0.4926,0.4706,0.4488,0.5648,0.5501,0.5362
+ViT-B-16-quickgelu,metaclip_400m,149.62,41.09,0.6962,0.6766,0.8960,0.9414,0.8570,0.9720,0.9870,0.3997,0.6536,0.7530,0.5648,0.7988,0.8714,0.5675,0.4984,0.5113,0.5242,0.5916,0.5853,0.5792
+ViT-L-14,commonpool_xl_clip_s13b_b90k,427.62,175.33,0.6956,0.6886,0.8996,0.9424,0.8650,0.9740,0.9920,0.4273,0.6761,0.7700,0.6040,0.8190,0.8858,0.5460,0.4833,0.4728,0.4622,0.5821,0.5691,0.5566
+convnext_base_w,laion2b_s13b_b82k_augreg,179.39,49.38,0.6922,0.7240,0.9144,0.9514,0.8730,0.9830,0.9920,0.4303,0.6786,0.7750,0.6140,0.8308,0.8964,0.5229,0.4504,0.4390,0.4276,0.5567,0.5494,0.5424
+EVA02-B-16,merged2b_s8b_b131k,149.69,41.09,0.6904,0.7146,0.9114,0.9468,0.8600,0.9660,0.9880,0.4215,0.6686,0.7629,0.5874,0.8056,0.8812,0.5325,0.4686,0.4722,0.4758,0.5522,0.5515,0.5509
+ViT-B-16,laion2b_s34b_b88k,149.62,41.09,0.6884,0.6984,0.9038,0.9456,0.8630,0.9790,0.9940,0.4231,0.6770,0.7706,0.5944,0.8178,0.8862,0.5220,0.4628,0.4601,0.4575,0.5549,0.5415,0.5288
+ViT-B-32-quickgelu,metaclip_fullcc,151.28,14.78,0.6879,0.6510,0.8766,0.9274,0.8080,0.9530,0.9730,0.3806,0.6411,0.7430,0.5518,0.7890,0.8650,0.5729,0.5351,0.5238,0.5126,0.5983,0.5883,0.5788
+convnext_base_w_320,laion_aesthetic_s13b_b82k_augreg,179.39,71.94,0.6870,0.7108,0.9032,0.9452,0.8910,0.9740,0.9900,0.4243,0.6722,0.7680,0.6040,0.8276,0.8912,0.5192,0.4599,0.4293,0.3989,0.5686,0.5475,0.5274
+ViT-B-16,datacomp_xl_s13b_b90k,149.62,41.09,0.6867,0.6756,0.8834,0.9300,0.8510,0.9720,0.9840,0.4016,0.6595,0.7562,0.5744,0.8086,0.8830,0.5408,0.4966,0.4832,0.4699,0.5730,0.5590,0.5457
+ViT-L-14,commonpool_xl_s13b_b90k,427.62,175.33,0.6817,0.6478,0.8656,0.9210,0.8200,0.9500,0.9740,0.3859,0.6310,0.7328,0.5446,0.7712,0.8484,0.5697,0.5115,0.5072,0.5030,0.5932,0.5894,0.5858
+ViT-B-32,laion2b_e16,151.28,14.78,0.6777,0.6638,0.8830,0.9322,0.8440,0.9650,0.9840,0.3913,0.6467,0.7481,0.5624,0.7956,0.8708,0.5343,0.4800,0.4602,0.4404,0.5669,0.5577,0.5490
+xlm-roberta-base-ViT-B-32,laion5b_s13b_b90k,366.12,105.87,0.6759,0.6448,0.8628,0.9168,0.8270,0.9640,0.9780,0.3778,0.6344,0.7348,0.5354,0.7798,0.8640,0.5519,0.4912,0.4827,0.4742,0.5834,0.5737,0.5645
+ViT-B-32,laion2b_s34b_b79k,151.28,14.78,0.6750,0.6678,0.8838,0.9310,0.8410,0.9620,0.9830,0.3934,0.6543,0.7561,0.5632,0.7984,0.8712,0.5254,0.4603,0.4479,0.4356,0.5603,0.5498,0.5398
+nllb-clip-large-siglip,v1,1195.5,1804.22,0.6738,0.7276,0.9224,0.9560,0.8330,0.9710,0.9910,0.4513,0.7084,0.8016,0.5386,0.7920,0.8712,0.4871,0.4138,0.4035,0.3933,0.5230,0.5135,0.5045
+ViT-L-14,laion400m_e32,427.62,175.33,0.6734,0.7022,0.9094,0.9458,0.8760,0.9780,0.9950,0.4300,0.6803,0.7740,0.5974,0.8218,0.8938,0.4815,0.4006,0.3932,0.3858,0.5245,0.5094,0.4950
+ViT-L-14,laion400m_e31,427.62,175.33,0.6728,0.7050,0.9068,0.9464,0.8720,0.9760,0.9950,0.4284,0.6797,0.7731,0.5974,0.8216,0.8934,0.4805,0.4022,0.3949,0.3877,0.5262,0.5075,0.4897
+ViT-B-32-256,datacomp_s34b_b86k,151.29,17.46,0.6712,0.6492,0.8722,0.9216,0.8480,0.9670,0.9850,0.3993,0.6543,0.7524,0.5792,0.8056,0.8810,0.5165,0.4554,0.4300,0.4048,0.5428,0.5437,0.5446
+ViT-B-32-quickgelu,metaclip_400m,151.28,14.78,0.6710,0.6234,0.8558,0.9152,0.7780,0.9350,0.9700,0.3591,0.6179,0.7209,0.5182,0.7640,0.8466,0.5657,0.5279,0.5097,0.4916,0.5907,0.5834,0.5764
+coca_ViT-B-32,laion2b_s13b_b90k,253.56,33.34,0.6689,0.6336,0.8566,0.9144,0.8160,0.9570,0.9740,0.3618,0.6136,0.7176,0.5456,0.7788,0.8556,0.5499,0.4924,0.4698,0.4472,0.5778,0.5751,0.5725
+ViT-B-16-plus-240,laion400m_e32,208.38,64.03,0.6676,0.6810,0.8888,0.9406,0.8650,0.9720,0.9880,0.4100,0.6619,0.7611,0.5858,0.8108,0.8806,0.4856,0.4146,0.4069,0.3993,0.5250,0.5104,0.4964
+ViT-B-16-plus-240,laion400m_e31,208.38,64.03,0.6671,0.6790,0.8878,0.9392,0.8560,0.9730,0.9860,0.4075,0.6604,0.7594,0.5834,0.8082,0.8796,0.4882,0.4206,0.4093,0.3980,0.5290,0.5131,0.4981
+ViT-L-14-336,openai,427.94,395.22,0.6666,0.6690,0.8900,0.9334,0.8770,0.9850,0.9940,0.3709,0.6162,0.7147,0.5794,0.8120,0.8792,0.5041,0.4201,0.4145,0.4089,0.5403,0.5324,0.5249
+roberta-ViT-B-32,laion2b_s12b_b32k,212.72,105.87,0.6641,0.6434,0.8674,0.9250,0.8170,0.9490,0.9720,0.3744,0.6309,0.7355,0.5436,0.7774,0.8568,0.5224,0.4602,0.4583,0.4564,0.5531,0.5426,0.5325
+ViT-B-16,laion400m_e32,149.62,41.09,0.6638,0.6566,0.8828,0.9300,0.8350,0.9680,0.9850,0.3835,0.6362,0.7386,0.5542,0.7964,0.8688,0.5033,0.4380,0.4326,0.4272,0.5294,0.5256,0.5219
+ViT-B-16,laion400m_e31,149.62,41.09,0.6621,0.6572,0.8808,0.9314,0.8330,0.9660,0.9840,0.3834,0.6366,0.7383,0.5514,0.7960,0.8688,0.5024,0.4255,0.4225,0.4195,0.5323,0.5276,0.5232
+ViT-B-32,datacomp_xl_s13b_b90k,151.28,14.78,0.6594,0.6108,0.8492,0.9092,0.7900,0.9390,0.9620,0.3714,0.6235,0.7268,0.5354,0.7778,0.8604,0.5317,0.4772,0.4594,0.4418,0.5589,0.5544,0.5502
+ViT-L-14,openai,427.62,175.33,0.6588,0.6496,0.8724,0.9204,0.8520,0.9740,0.9920,0.3651,0.6106,0.7113,0.5634,0.7938,0.8660,0.5047,0.4307,0.4136,0.3966,0.5457,0.5334,0.5218
+RN50x64,openai,623.26,552.65,0.6563,0.6898,0.8990,0.9432,0.8690,0.9820,0.9920,0.3524,0.5992,0.7033,0.5842,0.8046,0.8786,0.4779,0.4149,0.3936,0.3725,0.5280,0.5044,0.4820
+nllb-clip-base-siglip,v1,507.47,472.91,0.6545,0.6922,0.9002,0.9424,0.7920,0.9500,0.9820,0.4315,0.6913,0.7852,0.5116,0.7776,0.8628,0.4715,0.3776,0.3815,0.3854,0.5113,0.4999,0.4890
+convnext_base,laion400m_s13b_b51k,151.52,36.67,0.6541,0.6496,0.8814,0.9304,0.8380,0.9710,0.9910,0.3760,0.6315,0.7337,0.5470,0.7990,0.8676,0.4811,0.4146,0.4045,0.3944,0.5145,0.5052,0.4964
+RN50x16,openai,290.98,162.69,0.6518,0.6534,0.8710,0.9178,0.8570,0.9700,0.9880,0.3541,0.6002,0.7014,0.5536,0.7876,0.8670,0.4957,0.4311,0.3946,0.3584,0.5419,0.5275,0.5138
+ViT-B-16,openai,149.62,41.09,0.6507,0.6216,0.8572,0.9192,0.8220,0.9660,0.9900,0.3309,0.5842,0.6899,0.5242,0.7670,0.8462,0.5171,0.4487,0.4316,0.4146,0.5550,0.5441,0.5337
+ViT-B-32,laion400m_e31,151.28,14.78,0.6412,0.5970,0.8398,0.9036,0.7810,0.9380,0.9660,0.3420,0.6001,0.7059,0.5234,0.7634,0.8432,0.5059,0.4283,0.4262,0.4242,0.5477,0.5310,0.5152
+ViT-B-32,laion400m_e32,151.28,14.78,0.6412,0.5962,0.8396,0.9020,0.7770,0.9410,0.9680,0.3431,0.6000,0.7054,0.5244,0.7642,0.8454,0.5055,0.4272,0.4265,0.4258,0.5476,0.5304,0.5139
+ViT-B-32-quickgelu,laion400m_e32,151.28,14.78,0.6394,0.6170,0.8546,0.9086,0.7880,0.9400,0.9700,0.3533,0.6089,0.7165,0.5258,0.7672,0.8464,0.4884,0.4097,0.4072,0.4047,0.5280,0.5140,0.5006
+ViT-B-32-quickgelu,laion400m_e31,151.28,14.78,0.6389,0.6174,0.8548,0.9078,0.7870,0.9400,0.9730,0.3535,0.6100,0.7177,0.5254,0.7702,0.8490,0.4860,0.4054,0.4034,0.4015,0.5273,0.5120,0.4974
+RN50x4,openai,178.3,51.82,0.6373,0.6258,0.8476,0.9018,0.8210,0.9630,0.9830,0.3339,0.5812,0.6830,0.5296,0.7662,0.8490,0.4893,0.4104,0.3912,0.3720,0.5354,0.5202,0.5058
+nllb-clip-large,v1,1399.22,1468.46,0.6346,0.6090,0.8576,0.9202,0.7160,0.9220,0.9600,0.3617,0.6250,0.7304,0.4392,0.7086,0.8036,0.5097,0.4334,0.4299,0.4265,0.5395,0.5348,0.5304
+ViT-B-32,openai,151.28,14.78,0.6321,0.5878,0.8356,0.9002,0.7890,0.9490,0.9820,0.3044,0.5594,0.6687,0.5012,0.7500,0.8352,0.5054,0.4454,0.4125,0.3798,0.5492,0.5347,0.5210
+ViT-B-32-quickgelu,openai,151.28,14.78,0.6321,0.5878,0.8356,0.9002,0.7890,0.9490,0.9820,0.3044,0.5594,0.6687,0.5012,0.7500,0.8352,0.5054,0.4454,0.4125,0.3798,0.5492,0.5347,0.5210
+ViT-B-16,commonpool_l_laion_s1b_b8k,149.62,41.09,0.6274,0.5664,0.8114,0.8836,0.7230,0.9100,0.9510,0.3195,0.5756,0.6848,0.4652,0.7170,0.8088,0.5227,0.4558,0.4476,0.4395,0.5563,0.5463,0.5369
+ViT-B-16,datacomp_l_s1b_b8k,149.62,41.09,0.6267,0.5536,0.8090,0.8768,0.7320,0.9170,0.9480,0.3218,0.5747,0.6858,0.4872,0.7292,0.8246,0.5113,0.4613,0.4465,0.4318,0.5404,0.5317,0.5235
+RN101,openai,119.69,25.5,0.6249,0.5804,0.8228,0.8852,0.7900,0.9490,0.9740,0.3069,0.5546,0.6603,0.4982,0.7448,0.8250,0.4920,0.4347,0.4130,0.3913,0.5272,0.5170,0.5072
+RN101-quickgelu,openai,119.69,25.5,0.6249,0.5804,0.8228,0.8852,0.7900,0.9490,0.9740,0.3069,0.5546,0.6603,0.4982,0.7448,0.8250,0.4920,0.4347,0.4130,0.3913,0.5272,0.5170,0.5072
+ViT-B-16,commonpool_l_image_s1b_b8k,149.62,41.09,0.6164,0.5162,0.7908,0.8700,0.6890,0.8830,0.9270,0.2907,0.5449,0.6628,0.4338,0.6882,0.7902,0.5339,0.4932,0.4787,0.4643,0.5609,0.5513,0.5421
+ViT-B-16,commonpool_l_basic_s1b_b8k,149.62,41.09,0.6132,0.5250,0.7880,0.8670,0.6780,0.8780,0.9330,0.2862,0.5411,0.6535,0.4304,0.6844,0.7826,0.5285,0.4872,0.4812,0.4753,0.5517,0.5434,0.5355
+ViT-B-16,commonpool_l_text_s1b_b8k,149.62,41.09,0.6126,0.5336,0.7954,0.8650,0.6820,0.8880,0.9320,0.3020,0.5538,0.6674,0.4440,0.7012,0.7906,0.5147,0.4815,0.4570,0.4326,0.5394,0.5329,0.5266
+RN50,openai,102.01,18.18,0.6124,0.5736,0.8318,0.9004,0.8000,0.9500,0.9790,0.2854,0.5291,0.6459,0.4884,0.7278,0.8212,0.4766,0.3682,0.3622,0.3562,0.5318,0.5127,0.4945
+RN50-quickgelu,openai,102.01,18.18,0.6124,0.5736,0.8318,0.9004,0.8000,0.9500,0.9790,0.2854,0.5291,0.6459,0.4884,0.7278,0.8212,0.4766,0.3682,0.3622,0.3562,0.5318,0.5127,0.4945
+ViT-B-16,commonpool_l_clip_s1b_b8k,149.62,41.09,0.6054,0.5102,0.7764,0.8564,0.6810,0.8940,0.9430,0.2880,0.5368,0.6492,0.4436,0.6938,0.7954,0.5133,0.4493,0.4359,0.4225,0.5472,0.5377,0.5287
+ViT-B-16,commonpool_l_s1b_b8k,149.62,41.09,0.5635,0.4186,0.7040,0.7952,0.5710,0.8260,0.8940,0.2224,0.4565,0.5694,0.3486,0.6026,0.7096,0.5272,0.5029,0.4777,0.4526,0.5532,0.5428,0.5329
+nllb-clip-base,v1,501.89,369.6,0.5562,0.4740,0.7554,0.8350,0.5660,0.8210,0.8920,0.2649,0.5150,0.6321,0.3514,0.6100,0.7204,0.4708,0.3890,0.3904,0.3919,0.5085,0.4962,0.4844
+RN50-quickgelu,cc12m,102.01,18.18,0.5531,0.4696,0.7432,0.8284,0.6050,0.8600,0.9100,0.2412,0.4868,0.6024,0.3370,0.6136,0.7180,0.4644,0.4136,0.3876,0.3617,0.4970,0.4887,0.4807
+RN50,cc12m,102.01,18.18,0.5475,0.4594,0.7336,0.8234,0.6120,0.8540,0.9070,0.2367,0.4787,0.5947,0.3412,0.6072,0.7154,0.4553,0.4081,0.3828,0.3577,0.4918,0.4782,0.4653
+RN101-quickgelu,yfcc15m,119.69,25.5,0.4846,0.3474,0.6126,0.7192,0.5370,0.8030,0.8900,0.1779,0.3908,0.5043,0.2872,0.5456,0.6544,0.4226,0.3363,0.3147,0.2933,0.4723,0.4566,0.4416
+RN101,yfcc15m,119.69,25.5,0.4796,0.3402,0.6034,0.7124,0.5450,0.8010,0.8810,0.1772,0.3895,0.5040,0.3036,0.5536,0.6652,0.4088,0.3252,0.2987,0.2724,0.4643,0.4435,0.4239
+RN50-quickgelu,yfcc15m,102.01,18.18,0.4716,0.3180,0.5876,0.6968,0.4910,0.7840,0.8870,0.1645,0.3758,0.4891,0.2720,0.5260,0.6442,0.4209,0.3393,0.3130,0.2867,0.4634,0.4549,0.4469
+RN50,yfcc15m,102.01,18.18,0.4709,0.3092,0.5884,0.6952,0.5120,0.8040,0.8850,0.1645,0.3747,0.4894,0.2826,0.5350,0.6546,0.4106,0.3249,0.3023,0.2799,0.4536,0.4448,0.4364
+ViT-B-32,commonpool_m_image_s128m_b4k,151.28,14.78,0.3976,0.1896,0.4156,0.5312,0.2710,0.5300,0.6360,0.1054,0.2710,0.3703,0.1654,0.3758,0.4898,0.4804,0.4206,0.4026,0.3847,0.5261,0.5050,0.4849
+ViT-B-32,commonpool_m_clip_s128m_b4k,151.28,14.78,0.3923,0.1890,0.4106,0.5238,0.3040,0.5540,0.6590,0.1127,0.2718,0.3719,0.1826,0.3960,0.5070,0.4519,0.3685,0.3581,0.3478,0.4947,0.4815,0.4689
+ViT-B-32,commonpool_m_text_s128m_b4k,151.28,14.78,0.3917,0.1872,0.4200,0.5344,0.2900,0.5350,0.6590,0.1086,0.2712,0.3731,0.1818,0.3920,0.5008,0.4540,0.3863,0.3618,0.3373,0.5038,0.4831,0.4635
+ViT-B-32,commonpool_m_basic_s128m_b4k,151.28,14.78,0.3907,0.1874,0.4172,0.5350,0.2500,0.5190,0.6530,0.1007,0.2559,0.3578,0.1658,0.3528,0.4654,0.4727,0.4245,0.4015,0.3786,0.5182,0.4951,0.4732
+ViT-B-32,datacomp_m_s128m_b4k,151.28,14.78,0.3756,0.1812,0.4036,0.5188,0.2670,0.5200,0.6320,0.1100,0.2699,0.3750,0.1714,0.3770,0.4886,0.4337,0.3570,0.3287,0.3005,0.4914,0.4667,0.4433
+ViT-B-32,commonpool_m_laion_s128m_b4k,151.28,14.78,0.3741,0.1792,0.3964,0.5120,0.2760,0.5330,0.6220,0.1052,0.2606,0.3637,0.1654,0.3592,0.4660,0.4409,0.3585,0.3348,0.3113,0.4894,0.4743,0.4600
+ViT-B-32,commonpool_m_s128m_b4k,151.28,14.78,0.3449,0.1268,0.3164,0.4252,0.1930,0.4140,0.5360,0.0701,0.1941,0.2803,0.1246,0.3002,0.3996,0.4759,0.4354,0.3983,0.3614,0.5273,0.5003,0.4747
+ViT-B-32,commonpool_s_text_s13m_b4k,151.28,14.78,0.1922,0.0278,0.0950,0.1462,0.0570,0.1300,0.2000,0.0170,0.0551,0.0906,0.0258,0.0814,0.1234,0.3964,0.3312,0.3120,0.2929,0.4487,0.4231,0.3987
+ViT-B-32,commonpool_s_basic_s13m_b4k,151.28,14.78,0.1874,0.0256,0.0924,0.1466,0.0400,0.1220,0.2000,0.0144,0.0487,0.0782,0.0192,0.0646,0.1024,0.4008,0.3164,0.3033,0.2903,0.4526,0.4315,0.4115
+ViT-B-32,commonpool_s_clip_s13m_b4k,151.28,14.78,0.1857,0.0286,0.0938,0.1454,0.0400,0.1370,0.2070,0.0157,0.0515,0.0847,0.0292,0.0794,0.1206,0.3801,0.3282,0.2994,0.2708,0.4299,0.4055,0.3823
+ViT-B-32,commonpool_s_s13m_b4k,151.28,14.78,0.1789,0.0214,0.0690,0.1138,0.0300,0.1080,0.1760,0.0088,0.0351,0.0601,0.0176,0.0568,0.0902,0.3979,0.3338,0.3126,0.2914,0.4484,0.4249,0.4025
+ViT-B-32,commonpool_s_image_s13m_b4k,151.28,14.78,0.1536,0.0142,0.0594,0.0960,0.0200,0.0740,0.1160,0.0109,0.0385,0.0641,0.0150,0.0500,0.0842,0.3552,0.2669,0.2525,0.2382,0.4123,0.3875,0.3640
+ViT-B-32,datacomp_s_s13m_b4k,151.28,14.78,0.1536,0.0142,0.0594,0.0960,0.0200,0.0740,0.1160,0.0109,0.0385,0.0641,0.0150,0.0500,0.0842,0.3552,0.2669,0.2525,0.2382,0.4123,0.3875,0.3640
+ViT-B-32,commonpool_s_laion_s13m_b4k,151.28,14.78,0.1527,0.0176,0.0678,0.1056,0.0270,0.0870,0.1440,0.0102,0.0363,0.0607,0.0140,0.0442,0.0770,0.3463,0.2607,0.2410,0.2215,0.4049,0.3795,0.3554
+coca_ViT-B-32,mscoco_finetuned_laion2b_s13b_b90k,253.56,33.34,0.1306,0.0074,0.0214,0.0436,0.0110,0.0490,0.0990,0.0033,0.0137,0.0249,0.0088,0.0338,0.0552,0.3299,0.2484,0.2329,0.2175,0.3873,0.3604,0.3348

open_clip/docs/script_examples/clipa/vit_b16/i50_t16_finetune.sh

@@ -0,0 +1,27 @@
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion-400m' \
+    --dataset-type webdataset \
+    --lr "2.56e-5" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 3072 \
+    --wd 0.2 \
+    --batch-size 1024 \
+    --aug-cfg scale='(0.4, 1.0)' \
+    --epochs 1 \
+    --train-num-samples 131072000 \
+    --workers 6 \
+    --model ViT-B-16-CL16 \
+    --pretrained '/path/to/ckpt' \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --grad-checkpointing \
+    --log-every-n-steps 256 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipa/vit_b16/i50_t16_pretrain.sh

@@ -0,0 +1,26 @@
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion-400m' \
+    --dataset-type webdataset \
+    --lr "2.048e-3" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 782 \
+    --wd 0.2 \
+    --batch-size 8192 \
+    --aug-cfg scale='(0.4, 1.0)' \
+    --epochs 6 \
+    --workers 6 \
+    --model ViT-B-16-CL16 \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 112 \
+    --grad-checkpointing \
+    --log-every-n-steps 32 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipa/vit_l16/i17_t16_finetune.sh

@@ -0,0 +1,27 @@
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion-400m' \
+    --dataset-type webdataset \
+    --lr "2.24e-5" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 3571 \
+    --wd 0.2 \
+    --batch-size 896 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 1 \
+    --train-num-samples 131072000 \
+    --workers 6 \
+    --model ViT-L-16-CL16-GAP \
+    --pretrained '/path/to/ckpt' \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --grad-checkpointing \
+    --log-every-n-steps 293 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipa/vit_l16/i17_t16_pretrain.sh

@@ -0,0 +1,26 @@
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion-400m' \
+    --dataset-type webdataset \
+    --lr "1.024e-3" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 1563 \
+    --wd 0.2 \
+    --batch-size 4096 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 6 \
+    --workers 6 \
+    --model ViT-L-16-CL16-GAP \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 64 \
+    --grad-checkpointing \
+    --log-every-n-steps 64 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipa/vit_l16/i37_t8_finetune.sh

@@ -0,0 +1,27 @@
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion-400m' \
+    --dataset-type webdataset \
+    --lr "2.24e-5" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 3571 \
+    --wd 0.2 \
+    --batch-size 896 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 1 \
+    --train-num-samples 131072000 \
+    --workers 6 \
+    --model ViT-L-16-CL32-GAP \
+    --pretrained '/path/to/ckpt' \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --grad-checkpointing \
+    --log-every-n-steps 293 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipa/vit_l16/i37_t8_pretrain.sh

@@ -0,0 +1,26 @@
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion-400m' \
+    --dataset-type webdataset \
+    --lr "1.024e-3" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 1563 \
+    --wd 0.2 \
+    --batch-size 4096 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 6 \
+    --workers 6 \
+    --model ViT-L-16-CL8-Syntax-GAP \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 96 \
+    --grad-checkpointing \
+    --log-every-n-steps 64 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipav2/vit_h14/i257_t32_finetunex4.sh

@@ -0,0 +1,32 @@
+# have not been tested. use it at your own discretion
+# the original experiment was run on tpu v3-256.
+# this example script assumes 8 gpus, each with huge memory. Tune batchsize, warmup, and lr accordingly if you have different machine setups.
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion2b_or_datacomp1b' \
+    --train-num-samples 131072000 \
+    --dataset-type webdataset \
+    --lr "5.12e-5" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 800 \
+    --wd 0.2 \
+    --batch-size 4096 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 4 \
+    --workers 6 \
+    --model ViT-H-14-CL32-GAP \
+    --pretrained '/path/to/pretrain84_ckpt' \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 224 \
+    --force-patch-dropout 0.3 \
+    --grad-checkpointing \
+    --log-every-n-steps 64 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipav2/vit_h14/i50_t8_pretrain.sh

@@ -0,0 +1,30 @@
+# have not been tested. use it at your own discretion
+# the original experiment was run on tpu v3-256.
+# this example script assumes 8 gpus, each with huge memory. Tune batchsize, warmup, and lr accordingly if you have different machine setups.
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion2b_or_datacomp1b' \
+    --train-num-samples 4e8 \
+    --dataset-type webdataset \
+    --lr "2.048e-3" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 3200 \
+    --wd 0.2 \
+    --batch-size 8192 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 32 \
+    --workers 6 \
+    --model ViT-H-14-CL8-Syntax-GAP \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 84 \
+    --grad-checkpointing \
+    --log-every-n-steps 32 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/clipav2/vit_h14/i577_t32_finetunex1.sh

@@ -0,0 +1,32 @@
+# have not been tested. use it at your own discretion
+# the original experiment was run on tpu v3-256.
+# this example script assumes 8 gpus, each with huge memory. Tune batchsize, warmup, and lr accordingly if you have different machine setups.
+torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion2b_or_datacomp1b' \
+    --train-num-samples 131072000 \
+    --dataset-type webdataset \
+    --lr "6.4e-6" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 1600 \
+    --wd 0.2 \
+    --batch-size 2048 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs 1 \
+    --workers 6 \
+    --model ViT-H-14-CL32-GAP \
+    --pretrained '/path/to/finetune224_ckpt' \
+    --precision 'amp_bf16' \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 336 \
+    --force-patch-dropout 0.4 \
+    --grad-checkpointing \
+    --log-every-n-steps 64 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/imagenet/val'

open_clip/docs/script_examples/stability_example.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+#SBATCH --partition=g40423
+#SBATCH --job-name=testopenclip
+#SBATCH --nodes 30
+#SBATCH --ntasks-per-node=8
+#SBATCH --cpus-per-task=12
+#SBATCH --output=%x_%j.out
+#SBATCH --comment=laion
+#SBATCH --open-mode=append
+#SBATCH --exclusive
+
+module load openmpi
+module load cuda/11.7
+
+export MASTER_ADDR=`hostname`
+export MASTER_PORT=12802
+export NCCL_PROTO=simple
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1
+export NCCL_DEBUG=info
+
+export PYTHONFAULTHANDLER=1
+
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+cd /admin/home-mitchellw/open_clip/src
+export PYTHONPATH="$PYTHONPATH:/admin/home-mitchellw/open_clip/src"
+
+EXP_NAME="test-B-32-laion5b-lr1e-3-bs90k"
+
+srun --comment laion --cpu_bind=v --accel-bind=gn python -m open_clip_train.main \
+    --save-frequency 1 \
+    --train-data="pipe:aws s3 cp s3://s-datasets/laion5b/{laion2B-data/{000000..231349}.tar,laion2B-multi-data/{000000..226687}.tar,laion1B-nolang-data/{000000..127231}.tar} -" \
+    --train-num-samples 135646078 \
+    --dataset-type webdataset \
+    --dataset-resampled \
+    --warmup 2000 \
+    --batch-size=375 \
+    --epochs=97 \
+    --lr 1e-3 \
+    --workers=8 \
+    --report-to wandb \
+    --name ${EXP_NAME} \
+    --logs /scratch/logs/ \
+    --model ViT-B-32 \
+    --seed 0 \
+    --ddp-static-graph \
+    --local-loss \
+    --gather-with-grad \
+    --grad-checkpointing \
+    --precision amp_bfloat16 \
+    --wandb-project-name open_clip6 \
+    --resume "latest" \
+    --remote-sync s3://s-laion/mitchellw/logs

open_clip/scripts/clipav1_vit_l16_i37_t8.sh

@@ -0,0 +1,6 @@
+# eval on a single gpu
+CUDA_VISIBLE_DEVICES=2 TORCH_CUDNN_V8_API_ENABLED=1 TFDS_PREFETCH_SIZE=8192 python3 -m open_clip_train.main \
+    --model ViT-L-16-CL32-GAP \
+    --pretrained "/path/to/clipa_vit_l16_i37_t8.pt" \
+    --seed 0 \
+    --imagenet-val '/path/to/ImageNet/val'

open_clip/scripts/clipav2_vit_h14_i84_224_336_cl32_gap_datacomp1b.sh

@@ -0,0 +1,10 @@
+CUDA_VISIBLE_DEVICES=1 python3 -m open_clip_train.main \
+    --model ViT-H-14-CL32-GAP-BigVision \
+    --pretrained "/path/to/vit_h14_i84_224_336_cl32_gap_datacomp1b.pt" \
+    --force-image-size 336 \
+    --square-resize-only \
+    --interpolation 'bilinear' \
+    --image-mean 0.485 0.456 0.406 \
+    --image-std 0.229 0.224 0.225 \
+    --seed 0 \
+    --imagenet-val '/path/to/ImageNet/val'

open_clip/scripts/h14_84_8_pretrain.sh

@@ -0,0 +1,31 @@
+# 64k batchsize for 2.048e-3 lr
+TORCH_CUDNN_V8_API_ENABLED=1 torchrun --nproc_per_node 8 -m open_clip_train.main \
+    --save-frequency 1 \
+    --save-most-recent \
+    --zeroshot-frequency 1 \
+    --train-data '/path/to/laion' \
+    --dataset-type webdataset \
+    --lr "2.048e-3" \
+    --beta1 0.9 \
+    --beta2 0.95 \
+    --warmup 782 \
+    --wd 0.2 \
+    --batch-size 4096 \
+    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
+    --epochs=7 \
+    --workers=6 \
+    --model ViT-H-14-CL8-SyntaxMask-GAP \
+    --precision 'amp_bf16' \
+    --local-loss \
+    --gather-with-grad \
+    --force-image-size 84 \
+    --grad-checkpointing \
+    --log-every-n-steps 32 \
+    --seed 0 \
+    --logs ./logs/ \
+    --imagenet-val '/path/to/ImageNet/val' \
+    --name 'name' \
+    --report-to "wandb" \
+    --wandb-project-name "project_name"
+
+

open_clip/src/open_clip/constants.py

@@ -0,0 +1,11 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+INCEPTION_MEAN = (0.5, 0.5, 0.5)
+INCEPTION_STD = (0.5, 0.5, 0.5)
+
+# Default name for a weights file hosted on the Huggingface Hub.
+HF_WEIGHTS_NAME = "open_clip_pytorch_model.bin"  # default pytorch pkl
+HF_SAFE_WEIGHTS_NAME = "open_clip_model.safetensors"  # safetensors version
+HF_CONFIG_NAME = 'open_clip_config.json'

open_clip/src/open_clip/convert.py

@@ -0,0 +1,200 @@
+""" Conversion functions for 3rd part state-dicts and non-torch native checkpoint formats.
+"""
+from typing import Union
+
+import torch
+import numpy as np
+
+from .model import CLIP, CustomTextCLIP
+from .transformer import TextTransformer, Transformer
+
+
+@torch.no_grad()
+def load_big_vision_weights(model: CustomTextCLIP, checkpoint_path: str):
+    """ Load weights from .npz checkpoints for official Google big_vision image-text models
+
+    Currently the SigLIP source models are supported and a CustomTextCLIP destination model
+    w/ timm image encoder.
+    """
+    from timm.layers import resample_patch_embed, resample_abs_pos_embed
+
+    def _n2p(w, t=True, idx=None):
+        if idx is not None:
+            w = w[idx]
+        if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1:
+            w = w.flatten()
+        if t:
+            if w.ndim == 4:
+                w = w.transpose([3, 2, 0, 1])
+            elif w.ndim == 3:
+                w = w.transpose([2, 0, 1])
+            elif w.ndim == 2:
+                w = w.transpose([1, 0])
+        return torch.from_numpy(w)
+
+    w = np.load(checkpoint_path)
+    interpolation = 'bilinear'
+    antialias = False
+
+    def _convert_timm_img(module, prefix):
+        embed_conv_w = _n2p(w[f'{prefix}embedding/kernel'])
+        if embed_conv_w.shape[-2:] != module.patch_embed.proj.weight.shape[-2:]:
+            embed_conv_w = resample_patch_embed(
+                embed_conv_w,
+                module.patch_embed.proj.weight.shape[-2:],
+                interpolation=interpolation,
+                antialias=antialias,
+                verbose=True,
+            )
+        module.patch_embed.proj.weight.copy_(embed_conv_w)
+        module.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))
+
+        if module.cls_token is not None:
+            module.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
+
+        pos_embed_w = _n2p(w[f'{prefix}pos_embedding'], t=False)
+        if pos_embed_w.shape != module.pos_embed.shape:
+            assert False, f'{pos_embed_w.shape}, {module.pos_embed.shape}'
+            num_prefix_tokens = 0 if getattr(module, 'no_embed_class', False) else getattr(module, 'num_prefix_tokens', 1)
+            pos_embed_w = resample_abs_pos_embed(  # resize pos embedding when different size from pretrained weights
+                pos_embed_w,
+                new_size=module.patch_embed.grid_size,
+                num_prefix_tokens=num_prefix_tokens,
+                interpolation=interpolation,
+                antialias=antialias,
+                verbose=True,
+            )
+        module.pos_embed.copy_(pos_embed_w)
+
+        mha_sub, b_sub, ln1_sub = (0, 0, 1)
+        for i, block in enumerate(module.blocks.children()):
+            if f'{prefix}Transformer/encoderblock/LayerNorm_0/scale' in w:
+                block_prefix = f'{prefix}Transformer/encoderblock/'
+                idx = i
+            else:
+                block_prefix = f'{prefix}Transformer/encoderblock_{i}/'
+                idx = None
+            mha_prefix = block_prefix + f'MultiHeadDotProductAttention_{mha_sub}/'
+            block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale'], idx=idx))
+            block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias'], idx=idx))
+            block.attn.qkv.weight.copy_(torch.cat([
+                _n2p(w[f'{mha_prefix}{n}/kernel'], t=False, idx=idx).flatten(1).T for n in ('query', 'key', 'value')]))
+            block.attn.qkv.bias.copy_(torch.cat([
+                _n2p(w[f'{mha_prefix}{n}/bias'], t=False, idx=idx).reshape(-1) for n in ('query', 'key', 'value')]))
+            block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel'], idx=idx).flatten(1))
+            block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias'], idx=idx))
+            block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_{ln1_sub}/scale'], idx=idx))
+            block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_{ln1_sub}/bias'], idx=idx))
+            for r in range(2):
+                getattr(block.mlp, f'fc{r + 1}').weight.copy_(
+                    _n2p(w[f'{block_prefix}MlpBlock_{b_sub}/Dense_{r}/kernel'], idx=idx))
+                getattr(block.mlp, f'fc{r + 1}').bias.copy_(
+                    _n2p(w[f'{block_prefix}MlpBlock_{b_sub}/Dense_{r}/bias'], idx=idx))
+
+        module.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale']))
+        module.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias']))
+
+        if module.attn_pool is not None:
+            block_prefix = f'{prefix}MAPHead_0/'
+            mha_prefix = block_prefix + f'MultiHeadDotProductAttention_0/'
+            module.attn_pool.latent.copy_(_n2p(w[f'{block_prefix}probe'], t=False))
+            module.attn_pool.q.weight.copy_(_n2p(w[f'{mha_prefix}query/kernel'], t=False).flatten(1).T)
+            module.attn_pool.q.bias.copy_(_n2p(w[f'{mha_prefix}query/bias'], t=False).reshape(-1))
+            module.attn_pool.kv.weight.copy_(torch.cat([
+                _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('key', 'value')]))
+            module.attn_pool.kv.bias.copy_(torch.cat([
+                _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('key', 'value')]))
+            module.attn_pool.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))
+            module.attn_pool.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))
+            module.attn_pool.norm.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))
+            module.attn_pool.norm.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))
+            for r in range(2):
+                getattr(module.attn_pool.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_{r}/kernel']))
+                getattr(module.attn_pool.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_{r}/bias']))
+
+    def _convert_openclip_transformer(module: Transformer, prefix):
+        for i, block in enumerate(module.resblocks.children()):
+            block_prefix = f'{prefix}encoderblock_{i}/'
+            mha_prefix = block_prefix + f'MultiHeadDotProductAttention_0/'
+            block.ln_1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))
+            block.ln_1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))
+            block.attn.in_proj_weight.copy_(torch.cat([
+                _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')]))
+            block.attn.in_proj_bias.copy_(torch.cat([
+                _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')]))
+            block.attn.out_proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))
+            block.attn.out_proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))
+            block.ln_2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_1/scale']))
+            block.ln_2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_1/bias']))
+            block.mlp.c_fc.weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_0/kernel']))
+            block.mlp.c_fc.bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_0/bias']))
+            block.mlp.c_proj.weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_1/kernel']))
+            block.mlp.c_proj.bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_1/bias']))
+
+    def _convert_openclip_txt(module: TextTransformer, prefix):
+        module.token_embedding.weight.copy_(_n2p(w[f'{prefix}Embed_0/embedding'], t=False))
+        pos_embed_w = _n2p(w[f'{prefix}pos_embedding'], t=False).squeeze(0)
+        module.positional_embedding.copy_(pos_embed_w)
+        _convert_openclip_transformer(module.transformer, prefix=prefix + 'Encoder_0/')
+        module.ln_final.weight.copy_(_n2p(w[f'{prefix}Encoder_0/encoder_norm/scale']))
+        module.ln_final.bias.copy_(_n2p(w[f'{prefix}Encoder_0/encoder_norm/bias']))
+        if module.text_projection is not None:
+            module.text_projection.weight.copy_(_n2p(w[f'{prefix}head/kernel']))
+            module.text_projection.bias.copy_(_n2p(w[f'{prefix}head/bias']))
+
+    _convert_timm_img(model.visual.trunk, 'img/')
+    _convert_openclip_txt(model.text, 'txt/')
+    model.logit_bias.copy_(_n2p(w['b'])[0])
+    model.logit_scale.copy_(_n2p(w['t'])[0])
+
+
+@torch.no_grad()
+def convert_mobile_clip_state_dict(model: CustomTextCLIP, state_dict, fastvit = True):
+
+    def _convert_timm_img(state_dict):
+        if fastvit:
+            from timm.models.fastvit import checkpoint_filter_fn
+        else:
+            from timm.models.vision_transformer_hybrid import checkpoint_filter_fn
+        timm_state_dict = checkpoint_filter_fn(state_dict, model.visual.trunk)
+        timm_state_dict = {'visual.trunk.' + k: v for k, v in timm_state_dict.items()}
+        return timm_state_dict
+
+    def _convert_openclip_txt(state_dict, prefix='text_encoder.'):
+        text_dict = {}
+        for k, v in state_dict.items():
+            if not k.startswith(prefix):
+                continue
+            k = k.replace(prefix, '')
+            k = k.replace('projection_layer', 'text_projection')
+            k = k.replace('embedding_layer', 'token_embedding')
+            if k.startswith('positional_embedding.pos_embed.pos_embed'):
+                k = k.replace('positional_embedding.pos_embed.pos_embed', 'positional_embedding')
+                v = v.squeeze()
+            k = k.replace('final_layer_norm', 'ln_final')
+            k = k.replace('pre_norm_mha.0', 'ln_1')
+            k = k.replace('pre_norm_mha.1', 'attn')
+            k = k.replace('pre_norm_ffn.0', 'ln_2')
+            k = k.replace('pre_norm_ffn.1', 'mlp.c_fc')
+            k = k.replace('pre_norm_ffn.4', 'mlp.c_proj')
+            k = k.replace('qkv_proj.weight', 'in_proj_weight')
+            k = k.replace('qkv_proj.bias', 'in_proj_bias')
+            k = k.replace('transformer.', 'transformer.resblocks.')
+            text_dict['text.' + k] = v
+        return text_dict
+
+    image_dict = _convert_timm_img(state_dict)
+    text_dict = _convert_openclip_txt(state_dict)
+    out_dict = {**image_dict, **text_dict}
+    out_dict['logit_scale'] = state_dict['logit_scale']
+    return out_dict
+
+
+def convert_state_dict(model: Union[CustomTextCLIP, CLIP], state_dict):
+    if 'image_encoder.model.patch_embed.0.rbr_conv.0.conv.weight' in state_dict:
+        # Apple MobileCLIP s1 & s2 state_dicts (s0 and b not currently supported)
+        state_dict = convert_mobile_clip_state_dict(model, state_dict)
+    if 'image_encoder.model.patch_emb.0.block.conv.weight' in state_dict:
+        # convert b model
+        state_dict = convert_mobile_clip_state_dict(model, state_dict, fastvit=False)
+    return state_dict

open_clip/src/open_clip/factory.py

@@ -0,0 +1,578 @@
+import json
+import logging
+import os
+import re
+import warnings
+from copy import deepcopy
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+
+from .convert import convert_state_dict
+from .model import CLIP, CustomTextCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+    resize_pos_embed, get_cast_dtype, resize_text_pos_embed, set_model_preprocess_cfg
+from .coca_model import CoCa
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss, SigLipLoss
+from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained,\
+    list_pretrained_tags_by_model, download_pretrained_from_hf
+from .transform import image_transform_v2, AugmentationCfg, PreprocessCfg, merge_preprocess_dict, merge_preprocess_kwargs
+from .tokenizer import HFTokenizer, SimpleTokenizer, DEFAULT_CONTEXT_LENGTH
+
+HF_HUB_PREFIX = 'hf-hub:'
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = ('.json',)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f'*{ext}'))
+
+    for cf in config_files:
+        with open(cf, 'r') as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def list_models():
+    """ enumerate available model architectures based on config files """
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """ add model config path or file and update registry """
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    """ Fetch model config from builtin (local library) configs.
+    """
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None
+
+
+def _get_hf_config(
+        model_id: str,
+        cache_dir: Optional[str] = None,
+):
+    """ Fetch model config from HuggingFace Hub.
+    """
+    config_path = download_pretrained_from_hf(
+        model_id,
+        filename='open_clip_config.json',
+        cache_dir=cache_dir,
+    )
+    with open(config_path, 'r', encoding='utf-8') as f:
+        config = json.load(f)
+    return config
+
+
+def get_tokenizer(
+        model_name: str = '',
+        context_length: Optional[int] = None,
+        cache_dir: Optional[str] = None,
+        **kwargs,
+):
+    if model_name.startswith(HF_HUB_PREFIX):
+        model_name = model_name[len(HF_HUB_PREFIX):]
+        try:
+            config = _get_hf_config(model_name, cache_dir=cache_dir)['model_cfg']
+        except Exception:
+            tokenizer = HFTokenizer(
+                model_name,
+                context_length=context_length or DEFAULT_CONTEXT_LENGTH,
+                cache_dir=cache_dir,
+                **kwargs,
+            )
+            return tokenizer
+    else:
+        config = get_model_config(model_name)
+        assert config is not None, f"No valid model config found for {model_name}."
+
+    text_config = config.get('text_cfg', {})
+    if 'tokenizer_kwargs' in text_config:
+        tokenizer_kwargs = dict(text_config['tokenizer_kwargs'], **kwargs)
+    else:
+        tokenizer_kwargs = kwargs
+
+    if context_length is None:
+        context_length = text_config.get('context_length', DEFAULT_CONTEXT_LENGTH)
+
+    if 'hf_tokenizer_name' in text_config:
+        tokenizer = HFTokenizer(
+            text_config['hf_tokenizer_name'],
+            context_length=context_length,
+            cache_dir=cache_dir,
+            **tokenizer_kwargs,
+        )
+    else:
+        tokenizer = SimpleTokenizer(
+            context_length=context_length,
+            **tokenizer_kwargs,
+        )
+
+    return tokenizer
+
+
+def load_state_dict(
+        checkpoint_path: str,
+        device='cpu',
+        weights_only=True,
+):
+    # Check if safetensors or not and load weights accordingly
+    if str(checkpoint_path).endswith(".safetensors"):
+        from safetensors.torch import load_file
+        checkpoint = load_file(checkpoint_path, device=device)
+    else:
+        try:
+            checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=weights_only)
+        except TypeError:
+            checkpoint = torch.load(checkpoint_path, map_location=device)
+
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    elif isinstance(checkpoint, torch.jit.ScriptModule):
+        state_dict = checkpoint.state_dict()
+        for key in ["input_resolution", "context_length", "vocab_size"]:
+            state_dict.pop(key, None)
+    else:
+        state_dict = checkpoint
+    if next(iter(state_dict.items()))[0].startswith('module'):
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def load_checkpoint(
+        model: Union[CLIP, CustomTextCLIP],
+        checkpoint_path: str,
+        strict: bool = True,
+        weights_only: bool = True,
+        device='cpu',
+):
+    if Path(checkpoint_path).suffix in ('.npz', '.npy'):
+        # Separate path loading numpy big_vision (SigLIP) weights
+        from open_clip.convert import load_big_vision_weights
+        load_big_vision_weights(model, checkpoint_path)
+        return {}
+
+    state_dict = load_state_dict(checkpoint_path, device=device, weights_only=weights_only)
+
+    # Detect & convert 3rd party state_dicts -> open_clip
+    state_dict = convert_state_dict(model, state_dict)
+
+    # Detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+
+    # correct if logit_scale differs in being scaler vs 1d param
+    if 'logit_scale' in state_dict and model.logit_scale.ndim != state_dict['logit_scale'].ndim:
+        state_dict['logit_scale'] = state_dict['logit_scale'].reshape(model.logit_scale.shape)
+
+    # correct if logit_bias differs in being scaler vs 1d param
+    if 'logit_bias' in state_dict and model.logit_bias.ndim != state_dict['logit_bias'].ndim:
+        state_dict['logit_bias'] = state_dict['logit_bias'].reshape(model.logit_bias.shape)
+
+    # If loading a non-SigLIP model for SigLIP training. See https://github.com/mlfoundations/open_clip/issues/712
+    if 'logit_bias' not in state_dict and model.logit_bias is not None:
+        state_dict["logit_bias"] = torch.zeros_like(state_dict["logit_scale"])
+
+    # Certain text transformers no longer expect position_ids after transformers==4.31
+    position_id_key = 'text.transformer.embeddings.position_ids'
+    if position_id_key in state_dict and not hasattr(model, position_id_key):
+        del state_dict[position_id_key]
+
+    resize_pos_embed(state_dict, model)
+    resize_text_pos_embed(state_dict, model)
+
+    # Finally, load the massaged state_dict into model
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    return incompatible_keys
+
+
+def create_model(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        force_preprocess_cfg: Optional[Dict[str, Any]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        require_pretrained: bool = False,
+        load_weights_only: bool = True,
+        **model_kwargs,
+):
+    """Creates and configures a contrastive vision-language model.
+
+    Args:
+        model_name: Name of the model architecture to create. Can be a local model name
+            or a Hugging Face model ID prefixed with 'hf-hub:'.
+        pretrained: Tag/path for pretrained model weights. Can be:
+            - A pretrained tag name (e.g., 'openai')
+            - A path to local weights
+            - None to initialize with random weights
+        precision: Model precision/AMP configuration. Options:
+            - 'fp32': 32-bit floating point
+            - 'fp16'/'bf16': Mixed precision with FP32 for certain layers
+            - 'pure_fp16'/'pure_bf16': Pure 16-bit precision
+        device: Device to load the model on ('cpu', 'cuda', or torch.device object)
+        jit: If True, JIT compile the model
+        force_quick_gelu: Force use of QuickGELU activation
+        force_custom_text: Force use of custom text encoder
+        force_patch_dropout: Override default patch dropout value
+        force_image_size: Override default image size for vision encoder
+        force_preprocess_cfg: Override default preprocessing configuration
+        pretrained_image: Load pretrained weights for timm vision models
+        pretrained_hf: Load pretrained weights for HF text models when not loading CLIP weights
+        cache_dir: Override default cache directory for downloaded model files
+        output_dict: If True and model supports it, return dictionary of features
+        require_pretrained: Raise error if pretrained weights cannot be loaded
+        load_weights_only: Only deserialize model weights and unpickling torch checkpoints (for safety)
+        **model_kwargs: Additional keyword arguments passed to model constructor
+
+    Returns:
+        Created and configured model instance
+
+    Raises:
+        RuntimeError: If model config is not found or required pretrained weights
+            cannot be loaded
+
+    Examples:
+        # Create basic CLIP model
+        model = create_model('ViT-B/32')
+
+        # Create CLIP model with mixed precision on GPU
+        model = create_model('ViT-B/32', precision='fp16', device='cuda')
+
+        # Load pretrained OpenAI weights
+        model = create_model('ViT-B/32', pretrained='openai')
+
+        # Load Hugging Face model
+        model = create_model('hf-hub:organization/model-name')
+    """
+
+    force_preprocess_cfg = force_preprocess_cfg or {}
+    preprocess_cfg = asdict(PreprocessCfg())
+    has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
+    if has_hf_hub_prefix:
+        model_id = model_name[len(HF_HUB_PREFIX):]
+        checkpoint_path = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+        config = _get_hf_config(model_id, cache_dir=cache_dir)
+        preprocess_cfg = merge_preprocess_dict(preprocess_cfg, config['preprocess_cfg'])
+        model_cfg = config['model_cfg']
+        pretrained_hf = False  # override, no need to load original HF text weights
+    else:
+        model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+        checkpoint_path = None
+        model_cfg = None
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    model_cfg = model_cfg or get_model_config(model_name)
+    if model_cfg is not None:
+        logging.info(f'Loaded {model_name} model config.')
+    else:
+        logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
+        raise RuntimeError(f'Model config for {model_name} not found.')
+
+    if force_quick_gelu:
+        # override for use of QuickGELU on non-OpenAI transformer models
+        model_cfg["quick_gelu"] = True
+
+    if force_patch_dropout is not None:
+        # override the default patch dropout value
+        model_cfg["vision_cfg"]["patch_dropout"] = force_patch_dropout
+
+    if force_image_size is not None:
+        # override model config's image size
+        model_cfg["vision_cfg"]["image_size"] = force_image_size
+
+    is_timm_model = 'timm_model_name' in model_cfg.get('vision_cfg', {})
+    if pretrained_image:
+        if is_timm_model:
+            # pretrained weight loading for timm models set via vision_cfg
+            model_cfg['vision_cfg']['timm_model_pretrained'] = True
+        else:
+            assert False, 'pretrained image towers currently only supported for timm models'
+
+    # cast_dtype set for fp16 and bf16 (manual mixed-precision), not set for 'amp' or 'pure' modes
+    cast_dtype = get_cast_dtype(precision)
+    is_hf_model = 'hf_model_name' in model_cfg.get('text_cfg', {})
+    if is_hf_model:
+        # load pretrained weights for HF text model IFF no CLIP weights being loaded
+        model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf and not pretrained
+    custom_text = model_cfg.pop('custom_text', False) or force_custom_text or is_hf_model
+
+    model_cfg = dict(model_cfg, **model_kwargs)  # merge cfg dict w/ kwargs (kwargs overrides cfg)
+    if custom_text:
+        if "multimodal_cfg" in model_cfg:
+            model = CoCa(**model_cfg, cast_dtype=cast_dtype)
+        else:
+            model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
+    else:
+        model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
+    if precision in ("fp16", "bf16"):
+        dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+        # manual mixed precision that matches original OpenAI behaviour
+        if is_timm_model:
+            # FIXME this is a bit janky, create timm based model in low-precision and
+            # then cast only LayerNormFp32 instances back to float32 so they don't break.
+            # Why? The convert_weights_to_lp fn only works with native models.
+            model.to(device=device, dtype=dtype)
+            from .transformer import LayerNormFp32
+
+            def _convert_ln(m):
+                if isinstance(m, LayerNormFp32):
+                    m.weight.data = m.weight.data.to(torch.float32)
+                    m.bias.data = m.bias.data.to(torch.float32)
+            model.apply(_convert_ln)
+        else:
+            model.to(device=device)
+            convert_weights_to_lp(model, dtype=dtype)
+    elif precision in ("pure_fp16", "pure_bf16"):
+        dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+        model.to(device=device, dtype=dtype)
+    else:
+        model.to(device=device)
+
+    pretrained_loaded = False
+    if pretrained:
+        checkpoint_path = ''
+        pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+        if pretrained_cfg:
+            checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+            preprocess_cfg = merge_preprocess_dict(preprocess_cfg, pretrained_cfg)
+            pretrained_quick_gelu = pretrained_cfg.get('quick_gelu', False)
+            model_quick_gelu = model_cfg.get('quick_gelu', False)
+            if pretrained_quick_gelu and not model_quick_gelu:
+                warnings.warn(
+                    f'These pretrained weights were trained with QuickGELU activation but the model config does '
+                    f'not have that enabled. Consider using a model config with a "-quickgelu" suffix or enable with a flag.')
+            elif not pretrained_quick_gelu and model_quick_gelu:
+                warnings.warn(
+                    f'The pretrained weights were not trained with QuickGELU but this activation is enabled in the '
+                    f'model config, consider using a model config without QuickGELU or disable override flags.')
+        elif os.path.exists(pretrained):
+            checkpoint_path = pretrained
+
+        if checkpoint_path:
+            logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+            load_checkpoint(model, checkpoint_path, weights_only=load_weights_only)
+        else:
+            error_str = (
+                f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+                f' Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+            logging.warning(error_str)
+            raise RuntimeError(error_str)
+        pretrained_loaded = True
+    elif has_hf_hub_prefix:
+        logging.info(f'Loading pretrained {model_name} weights ({checkpoint_path}).')
+        load_checkpoint(model, checkpoint_path, weights_only=load_weights_only)
+        pretrained_loaded = True
+
+    if require_pretrained and not pretrained_loaded:
+        # callers of create_model_from_pretrained always expect pretrained weights
+        raise RuntimeError(
+            f'Pretrained weights were required for (model: {model_name}, pretrained: {pretrained}) but not loaded.')
+
+    if output_dict and hasattr(model, "output_dict"):
+        model.output_dict = True
+
+    if jit:
+        model = torch.jit.script(model)
+
+    # set image preprocessing configuration in model attributes for convenience
+    if getattr(model.visual, 'image_size', None) is not None:
+        # use image_size set on model creation (via config or force_image_size arg)
+        force_preprocess_cfg['size'] = model.visual.image_size
+    set_model_preprocess_cfg(model, merge_preprocess_dict(preprocess_cfg, force_preprocess_cfg))
+
+    return model
+
+
+def create_loss(args):
+    if args.distill:
+        return DistillClipLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    elif "coca" in args.model.lower():
+        return CoCaLoss(
+            caption_loss_weight=args.coca_caption_loss_weight,
+            clip_loss_weight=args.coca_contrastive_loss_weight,
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    elif args.siglip:
+        assert not args.horovod, "Horovod not currently supported for SigLip"
+        return SigLipLoss(
+            rank=args.rank,
+            world_size=args.world_size,
+            dist_impl=args.loss_dist_impl,  # siglip has multiple distributed implementations to choose from
+        )
+
+    return ClipLoss(
+        local_loss=args.local_loss,
+        gather_with_grad=args.gather_with_grad,
+        cache_labels=True,
+        rank=args.rank,
+        world_size=args.world_size,
+        use_horovod=args.horovod,
+    )
+
+
+def create_model_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        image_interpolation: Optional[str] = None,
+        image_resize_mode: Optional[str] = None,  # only effective for inference
+        aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        load_weights_only: bool = True,
+        **model_kwargs,
+):
+    force_preprocess_cfg = merge_preprocess_kwargs(
+        {},
+        mean=image_mean,
+        std=image_std,
+        interpolation=image_interpolation,
+        resize_mode=image_resize_mode,
+    )
+
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_patch_dropout=force_patch_dropout,
+        force_image_size=force_image_size,
+        force_preprocess_cfg=force_preprocess_cfg,
+        pretrained_image=pretrained_image,
+        pretrained_hf=pretrained_hf,
+        cache_dir=cache_dir,
+        output_dict=output_dict,
+        load_weights_only=load_weights_only,
+        **model_kwargs,
+    )
+
+    pp_cfg = PreprocessCfg(**model.visual.preprocess_cfg)
+
+    preprocess_train = image_transform_v2(
+        pp_cfg,
+        is_train=True,
+        aug_cfg=aug_cfg,
+    )
+    preprocess_val = image_transform_v2(
+        pp_cfg,
+        is_train=False,
+    )
+
+    return model, preprocess_train, preprocess_val
+
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        image_interpolation: Optional[str] = None,
+        image_resize_mode: Optional[str] = None,  # only effective for inference
+        return_transform: bool = True,
+        cache_dir: Optional[str] = None,
+        load_weights_only: bool = True,
+        **model_kwargs,
+):
+    force_preprocess_cfg = merge_preprocess_kwargs(
+        {},
+        mean=image_mean,
+        std=image_std,
+        interpolation=image_interpolation,
+        resize_mode=image_resize_mode,
+    )
+
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_image_size=force_image_size,
+        force_preprocess_cfg=force_preprocess_cfg,
+        cache_dir=cache_dir,
+        require_pretrained=True,
+        load_weights_only=load_weights_only,
+        **model_kwargs,
+    )
+
+    if not return_transform:
+        return model
+
+    preprocess = image_transform_v2(
+        PreprocessCfg(**model.visual.preprocess_cfg),
+        is_train=False,
+    )
+
+    return model, preprocess

open_clip/src/open_clip/hf_model.py

@@ -0,0 +1,193 @@
+""" huggingface model adapter
+
+Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
+"""
+import re
+
+import torch
+import torch.nn as nn
+from torch import TensorType
+
+try:
+    import transformers
+    from transformers import AutoModel, AutoTokenizer, AutoConfig, PretrainedConfig
+    from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
+        BaseModelOutputWithPoolingAndCrossAttentions
+except ImportError as e:
+    transformers = None
+
+
+    class BaseModelOutput:
+        pass
+
+
+    class PretrainedConfig:
+        pass
+
+from .hf_configs import arch_dict
+
+
+# utils
+def _camel2snake(s):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
+
+
+# TODO: ?last - for gpt-like models
+_POOLERS = {}
+
+
+def register_pooler(cls):
+    """Decorator registering pooler class"""
+    _POOLERS[_camel2snake(cls.__name__)] = cls
+    return cls
+
+
+@register_pooler
+class MeanPooler(nn.Module):
+    """Mean pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
+        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
+
+
+@register_pooler
+class MaxPooler(nn.Module):
+    """Max pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state.masked_fill(attention_mask.unsqueeze(-1), -torch.inf)
+        return masked_output.max(1).values
+
+
+@register_pooler
+class ClsPooler(nn.Module):
+    """CLS token pooling"""
+
+    def __init__(self, use_pooler_output=True):
+        super().__init__()
+        self.cls_token_position = 0
+        self.use_pooler_output = use_pooler_output
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        if (self.use_pooler_output and
+            isinstance(x, (BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions)) and
+            (x.pooler_output is not None)
+        ):
+            return x.pooler_output
+
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+@register_pooler
+class ClsLastHiddenStatePooler(nn.Module):
+    """CLS token pooling
+    NOTE: this is equivalent to ClsPooler above with use_pooler_output=False
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.cls_token_position = 0
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+class HFTextEncoder(nn.Module):
+    """HuggingFace model adapter"""
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            model_name_or_path: str,
+            output_dim: int,
+            config: PretrainedConfig = None,
+            pooler_type: str = None,
+            proj_type: str = None,
+            pretrained: bool = True,
+            output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        self.output_dim = output_dim
+
+        # TODO: find better way to get this information
+        uses_transformer_pooler = (pooler_type == "cls_pooler")
+
+        if transformers is None:
+            raise RuntimeError("Please `pip install transformers` to use pre-trained HuggingFace models")
+        if config is None:
+            self.config = AutoConfig.from_pretrained(model_name_or_path)
+            create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else (
+                AutoModel.from_config, self.config)
+            # TODO: do all model configs have this attribute? PretrainedConfig does so yes??
+            if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder:
+                self.transformer = create_func(model_args)
+                self.transformer = self.transformer.encoder
+            else:
+                self.transformer = create_func(model_args, add_pooling_layer=uses_transformer_pooler)
+        else:
+            self.config = config
+            self.transformer = AutoModel.from_config(config)
+        if pooler_type is None:  # get default arch pooler
+            pooler_type = (arch_dict[self.config.model_type]["pooler"])
+
+        # FIXME downstream users of OpenCLIP models use these attr, need to verify valid across all models
+        self.vocab_size = getattr(self.config, 'vocab_size', 0)
+        self.context_length = getattr(self.config, 'max_position_embeddings', 0)
+
+        self.pooler = _POOLERS[pooler_type]()
+
+        d_model = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["width"])
+        if (d_model == output_dim) and (proj_type is None):  # do we always need a proj?
+            self.proj = nn.Identity()
+        elif proj_type == 'linear':
+            self.proj = nn.Linear(d_model, output_dim, bias=False)
+        elif proj_type == 'mlp':
+            hidden_size = (d_model + output_dim) // 2
+            self.proj = nn.Sequential(
+                nn.Linear(d_model, hidden_size, bias=False),
+                nn.GELU(),
+                nn.Linear(hidden_size, output_dim, bias=False),
+            )
+
+    def forward(self, x: TensorType):
+        attn_mask = (x != self.config.pad_token_id).long()
+        out = self.transformer(input_ids=x, attention_mask=attn_mask)
+        pooled_out = self.pooler(out, attn_mask)
+        projected = self.proj(pooled_out)
+
+        seq_len = out.last_hidden_state.shape[1]
+        tokens = (
+            out.last_hidden_state[:, torch.arange(seq_len) != self.pooler.cls_token_position, :] 
+            if type(self.pooler) == ClsPooler 
+            else out.last_hidden_state
+        )
+        
+        if self.output_tokens:
+            return projected, tokens
+        return projected
+
+    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        if not unlocked_layers:  # full freezing
+            for n, p in self.transformer.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+            return
+
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
+        embeddings = getattr(
+            self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
+        modules = [embeddings, *layer_list][:-unlocked_layers]
+        # freeze layers
+        for module in modules:
+            for n, p in module.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.gradient_checkpointing_enable()
+
+    def init_parameters(self):
+        pass

open_clip/src/open_clip/model.py

@@ -0,0 +1,650 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import copy
+import logging
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from functools import partial
+
+from .hf_model import HFTextEncoder
+from .modified_resnet import ModifiedResNet
+from .timm_model import TimmModel
+from .transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer,\
+    text_global_pool
+from .utils import to_2tuple
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer (overrides pool_type)
+    attn_pooler_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    no_ln_pre: bool = False  # disable pre transformer LayerNorm
+    pos_embed_type: str = 'learnable'
+    final_ln_after_pool: bool = False  # apply final LayerNorm after pooling
+    pool_type: str = 'tok'
+    output_tokens: bool = False
+    act_kwargs: Optional[dict] = None
+    norm_kwargs: Optional[dict] = None
+
+    timm_model_name: Optional[str] = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    hf_tokenizer_name: Optional[str] = None
+    tokenizer_kwargs: Optional[dict] = None
+
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    mlp_ratio: float = 4.0
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    embed_cls: bool = False
+    pad_id: int = 0
+    no_causal_mask: bool = False  # disable causal masking
+    final_ln_after_pool: bool = False  # apply final LayerNorm after pooling
+    pool_type: str = 'argmax'
+    proj_bias: bool = False
+    proj_type: str = 'linear'  # control final text projection, 'none' forces no projection
+    output_tokens: bool = False
+    act_kwargs: dict = None
+    norm_kwargs: dict = None
+
+    # HuggingFace specific text tower config
+    hf_model_name: Optional[str] = None
+    hf_model_pretrained: bool = True
+    hf_proj_type: str = 'mlp'
+    hf_pooler_type: str = 'mean_pooler'  # attentional pooling for HF models
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        if vision_cfg.norm_kwargs:
+            norm_layer = partial(norm_layer, **vision_cfg.norm_kwargs)
+        if vision_cfg.act_kwargs is not None:
+            act_layer = partial(act_layer, **vision_cfg.act_kwargs)
+
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            attentional_pool=vision_cfg.attentional_pool,
+            attn_pooler_queries=vision_cfg.attn_pooler_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            pos_embed_type=vision_cfg.pos_embed_type,
+            no_ln_pre=vision_cfg.no_ln_pre,
+            final_ln_after_pool=vision_cfg.final_ln_after_pool,
+            pool_type=vision_cfg.pool_type,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj_type=text_cfg.hf_proj_type,
+            pooler_type=text_cfg.hf_pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        if text_cfg.norm_kwargs:
+            norm_layer = partial(norm_layer, **text_cfg.norm_kwargs)
+        if text_cfg.act_kwargs is not None:
+            act_layer = partial(act_layer, **text_cfg.act_kwargs)
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            mlp_ratio=text_cfg.mlp_ratio,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            no_causal_mask=text_cfg.no_causal_mask,
+            pad_id=text_cfg.pad_id,
+            pool_type=text_cfg.pool_type,
+            proj_type=text_cfg.proj_type,
+            proj_bias=text_cfg.proj_bias,
+            output_tokens=text_cfg.output_tokens,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            init_logit_scale: float = np.log(1 / 0.07),
+            init_logit_bias: Optional[float] = None,
+            nonscalar_logit_scale: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.context_length = text.context_length
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.text_pool_type = text.pool_type
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        lshape = [1] if nonscalar_logit_scale else []
+        self.logit_scale = nn.Parameter(torch.ones(lshape) * init_logit_scale)
+        if init_logit_bias is not None:
+            self.logit_bias = nn.Parameter(torch.ones(lshape) * init_logit_bias)
+        else:
+            self.logit_bias = None
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # for timm optimizers, 1d params like logit_scale, logit_bias, ln/bn scale, biases are excluded by default
+        no_wd = {'positional_embedding'}
+        if hasattr(self.visual, 'no_weight_decay'):
+            for n in self.visual.no_weight_decay():
+                no_wd.add('visual.' + n)
+        return no_wd
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        x, _ = text_global_pool(x, text, self.text_pool_type)
+        if self.text_projection is not None:
+            if isinstance(self.text_projection, nn.Linear):
+                x = self.text_projection(x)
+            else:
+                x = x @ self.text_projection
+
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def get_logits(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        image_logits = self.logit_scale.exp() * image_features @ text_features.T
+        if self.logit_bias is not None:
+            image_logits += self.logit_bias
+        text_logits = image_logits.T
+        return image_logits, text_logits
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+
+        if self.output_dict:
+            out_dict = {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+            if self.logit_bias is not None:
+                out_dict['logit_bias'] = self.logit_bias
+            return out_dict
+
+        if self.logit_bias is not None:
+            return image_features, text_features, self.logit_scale.exp(), self.logit_bias
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            init_logit_scale: float = np.log(1 / 0.07),
+            init_logit_bias: Optional[float] = None,
+            nonscalar_logit_scale: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+
+        lshape = [1] if nonscalar_logit_scale else []
+        self.logit_scale = nn.Parameter(torch.ones(lshape) * init_logit_scale)
+        if init_logit_bias is not None:
+            self.logit_bias = nn.Parameter(torch.ones(lshape) * init_logit_bias)
+        else:
+            self.logit_bias = None
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # for timm optimizers, 1d params like logit_scale, logit_bias, ln/bn scale, biases are excluded by default
+        no_wd = set()
+        if hasattr(self.visual, 'no_weight_decay'):
+            for n in self.visual.no_weight_decay():
+                no_wd.add('visual.' + n)
+        if hasattr(self.text, 'no_weight_decay'):
+            for n in self.visual.no_weight_decay():
+                no_wd.add('text.' + n)
+        return no_wd
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def get_logits(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        image_logits = self.logit_scale.exp() * image_features @ text_features.T
+        if self.logit_bias is not None:
+            image_logits += self.logit_bias
+        text_logits = image_logits.T
+        return image_logits, text_logits
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+
+        if self.output_dict:
+            out_dict = {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+            if self.logit_bias is not None:
+                out_dict['logit_bias'] = self.logit_bias
+            return out_dict
+
+        if self.logit_bias is not None:
+            return image_features, text_features, self.logit_scale.exp(), self.logit_bias
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed
+
+
+def resize_text_pos_embed(state_dict, model, interpolation: str = 'linear', antialias: bool = False):
+    old_pos_embed = state_dict.get('positional_embedding', None)
+    if old_pos_embed is None:
+        return
+    # FIXME add support for text cls_token
+    model_pos_embed = getattr(model, 'positional_embedding', None)
+    if model_pos_embed is None:
+        model_pos_embed = getattr(model.text, 'positional_embedding', None)
+
+    old_num_pos = old_pos_embed.shape[0]
+    old_width = old_pos_embed.shape[1]
+    num_pos = model_pos_embed.shape[0]
+    width = model_pos_embed.shape[1]
+    assert old_width == width, 'text pos_embed width changed!'
+    if old_num_pos == num_pos:
+        return
+
+    logging.info('Resizing text position embedding num_pos from %s to %s', old_num_pos, num_pos)
+    old_pos_embed = old_pos_embed.reshape(1, old_num_pos, old_width).permute(0, 2, 1)
+    old_pos_embed = F.interpolate(
+        old_pos_embed,
+        size=num_pos,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    old_pos_embed = old_pos_embed.permute(0, 2, 1)[0]
+    new_pos_embed = old_pos_embed
+
+    state_dict['positional_embedding'] = new_pos_embed
+
+
+def get_model_preprocess_cfg(model):
+    module = getattr(model, 'visual', model)
+    preprocess_cfg = getattr(module, 'preprocess_cfg', {})
+    if not preprocess_cfg:
+        # use separate legacy attributes if preprocess_cfg dict not found
+        size = getattr(module, 'image_size')
+        if size is not None:
+            preprocess_cfg['size'] = size
+        mean = getattr(module, 'image_mean', None)
+        if mean is not None:
+            preprocess_cfg['mean'] = mean
+        std = getattr(module, 'image_std', None)
+        if std is not None:
+            preprocess_cfg['std'] = std
+    return preprocess_cfg
+
+
+def set_model_preprocess_cfg(model, preprocess_cfg: Dict[str, Any]):
+    module = getattr(model, 'visual', model)
+    module.image_mean = preprocess_cfg['mean']  # legacy attribute, keeping for bwd compat
+    module.image_std = preprocess_cfg['std']  # legacy attribute, keeping for bwd compat
+    module.preprocess_cfg = copy.deepcopy(preprocess_cfg)  # new attr, package all pp cfg as dict
+
+
+def get_model_tokenize_cfg(model):
+    module = getattr(model, 'text', model)
+    cfg = {}
+    context_length = getattr(module, 'context_length', None)
+    if context_length is not None:
+        cfg['context_length'] = context_length
+    vocab_size = getattr(module, 'vocab_size', None)
+    if vocab_size is not None:
+        cfg['vocab_size'] = vocab_size
+    return cfg

open_clip/src/open_clip/model_configs/EVA01-g-14-plus.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva_giant_patch14_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    },
+    "custom_text": true
+}

open_clip/src/open_clip/model_configs/EVA02-B-16.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_base_patch16_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    },
+    "custom_text": true
+}

open_clip/src/open_clip/model_configs/EVA02-E-14-plus.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_enormous_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32
+    },
+    "custom_text": true
+}

open_clip/src/open_clip/model_configs/EVA02-L-14-336.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "timm_model_name": "eva02_large_patch14_clip_336",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

open_clip/src/open_clip/model_configs/RN50x16-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 768,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 384,
+        "layers": [
+            6,
+            8,
+            18,
+            8
+        ],
+        "width": 96,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

open_clip/src/open_clip/model_configs/RN50x4-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 640,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 288,
+        "layers": [
+            4,
+            6,
+            10,
+            6
+        ],
+        "width": 80,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

open_clip/src/open_clip/model_configs/RN50x4.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 288,
+        "layers": [
+            4,
+            6,
+            10,
+            6
+        ],
+        "width": 80,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

open_clip/src/open_clip/model_configs/RN50x64.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": [
+            3,
+            15,
+            36,
+            10
+        ],
+        "width": 128,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 12
+    }
+}

open_clip/src/open_clip/model_configs/ViT-B-16-SigLIP-256.json

@@ -0,0 +1,29 @@
+{
+    "embed_dim": 768,
+    "init_logit_bias": -10,
+    "custom_text": true,
+    "vision_cfg": {
+        "image_size": 256,
+        "timm_model_name": "vit_base_patch16_siglip_256",
+        "timm_model_pretrained": false,
+        "timm_pool": "map",
+        "timm_proj": "none"
+    },
+    "text_cfg": {
+        "context_length": 64,
+        "vocab_size": 32000,
+        "hf_tokenizer_name": "timm/ViT-B-16-SigLIP",
+        "tokenizer_kwargs": {
+            "clean": "canonicalize"
+        },
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "no_causal_mask": true,
+        "proj_bias": true,
+        "pool_type": "last",
+        "norm_kwargs":{
+            "eps": 1e-6
+        }
+    }
+}

open_clip/src/open_clip/model_configs/ViT-B-16-plus-240.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 240,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

open_clip/src/open_clip/model_configs/ViT-B-16-plus.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

open_clip/src/open_clip/model_configs/ViT-L-14-CLIPA-336.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14,
+        "no_ln_pre": true,
+        "pool_type": "avg",
+        "final_ln_after_pool": true
+    },
+    "text_cfg": {
+        "context_length": 32,
+        "vocab_size": 32000,
+        "hf_tokenizer_name": "bert-base-uncased",
+        "tokenizer_kwargs": {
+            "strip_sep_token": true
+        },
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "pool_type": "last",
+        "no_causal_mask": true
+    }
+}

open_clip/src/open_clip/model_configs/ViT-L-14-CLIPA.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14,
+        "no_ln_pre": true,
+        "pool_type": "avg",
+        "final_ln_after_pool": true
+    },
+    "text_cfg": {
+        "context_length": 32,
+        "vocab_size": 32000,
+        "hf_tokenizer_name": "bert-base-uncased",
+        "tokenizer_kwargs": {
+            "strip_sep_token": true
+        },
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "pool_type": "last",
+        "no_causal_mask": true
+    }
+}