add norm onnx

Author: xxaier

onnx/README.md

@@ -30,7 +30,7 @@ ONNX(Open Neural Network Exchange)，开放神经网络交换，用于在各种
 
 * `./bash.sh` 在本地进入容器的 bash，方便调试
 
-* `./export.sh` 运行容器，导出 onnx
+* `./export.sh` 运行容器，下载 pytorch 模型，然后转换为 onnx
 
     设置环境变量 MODEL ，可以配置导出、测试脚本运行的模型 。
 
@@ -39,6 +39,17 @@ ONNX(Open Neural Network Exchange)，开放神经网络交换，用于在各种
     * AltCLIP-XLMR-L
     * AltCLIP-XLMR-L-m9
 
+    运行后将生成 4 个 onnx 文件和很多权重文件
+
+    * onnx/AltCLIP-XLMR-L-m18/onnx/Img.onnx
+    * onnx/AltCLIP-XLMR-L-m18/onnx/ImgNorm.onnx
+    * onnx/AltCLIP-XLMR-L-m18/onnx/Txt.onnx
+    * onnx/AltCLIP-XLMR-L-m18/onnx/TxtNorm.onnx
+
+    其中 Norm 代表输出归一化的向量，如果想把生成的文本向量和图片向量存入向量数据库，进行相似性搜索，请用归一化的向量。
+
+    具体用见下文的 onnx 模型的测试脚本。
+
 * `./dist.sh` 运行容器，导出以上 3 个模型的 onnx，并打包放到 dist 目录下。
 
 ### 目录
@@ -67,17 +78,10 @@ onnxruntime 有很多版本可以选择，见[onnxruntime](https://onnxruntime.a
 
 请先安装 [direnv](https://github.com/direnv/direnv/blob/master/README.md) 并在本目录下 `direnv allow` 或者手工 `source .envrc` 来设置 PYTHONPATH 环境变量。
 
-* [./test/onnx/onnx_img.py](./test/onnx/onnx_img.py)  生成图片向量
+* [./test/onnx/onnx_img.py](./test/onnx/onnx_img.py)  生成图片向量 (norm 代表归一化的向量，可用于向量搜索)
 * [./test/onnx/onnx_txt.py](./test/onnx/onnx_txt.py)  生成文本向量
 * [./test/onnx/onnx_test.py](./test/onnx/onnx_test.py) 匹配图片向量和文本向量，进行零样本分类
 
-如果想把生成的文本向量和图片向量存入数据库，进行相似性搜索，请先对特征进行归一化。
-
-```python
-image_features /= image_features.norm(dim=-1, keepdim=True)
-text_features /= text_features.norm(dim=-1, keepdim=True)
-```
-
 可借助向量数据库，提升零样本分类的准确性，参见[ECCV 2022 | 无需下游训练，Tip-Adapter 大幅提升 CLIP 图像分类准确率](https://cloud.tencent.com/developer/article/2126102)。
 
 #### pytorch 模型

onnx/export/onnx_export.py

@@ -3,7 +3,7 @@
 from PIL import Image
 from os import makedirs
 from os.path import join
-from misc.clip_model import TXT, IMG
+from misc.clip_model import TXT, IMG, TXT_NORM, IMG_NORM
 from misc.config import ONNX_FP, opset_version, IMG_DIR
 from misc.proc import transform, tokenizer
 import torch
@@ -17,8 +17,8 @@
 image = torch.tensor(image)
 
 
-def onnx_export(outdir, model, args, **kwds):
-  name = f'{outdir}.onnx'
+def onnx_export(model, args, **kwds):
+  name = f'{model.__class__.__name__}.onnx'
   fp = join(ONNX_FP, name)
   torch.onnx.export(
       model,
@@ -35,26 +35,28 @@ def onnx_export(outdir, model, args, **kwds):
 
 
 # 参考 https://github.com/OFA-Sys/Chinese-CLIP/blob/master/cn_clip/deploy/pytorch_to_onnx.py
-
-onnx_export('txt',
-            TXT,
-            tokenizer(['a photo of cat', 'a image of cat'], ),
-            input_names=['input', 'attention_mask'],
-            dynamic_axes={
-                'input': {
-                    0: 'batch',
-                    1: 'batch',
-                },
-                'attention_mask': {
-                    0: 'batch',
-                    1: 'batch',
-                }
-            })
-
-onnx_export('img',
-            IMG,
-            image,
-            input_names=['input'],
-            dynamic_axes={'input': {
-                0: 'batch'
-            }})
+def export(txt, img):
+  onnx_export(txt,
+              tokenizer(['a photo of cat', 'a image of cat'], ),
+              input_names=['input', 'attention_mask'],
+              dynamic_axes={
+                  'input': {
+                      0: 'batch',
+                      1: 'batch',
+                  },
+                  'attention_mask': {
+                      0: 'batch',
+                      1: 'batch',
+                  }
+              })
+
+  onnx_export(img,
+              image,
+              input_names=['input'],
+              dynamic_axes={'input': {
+                  0: 'batch'
+              }})
+
+
+export(TXT, IMG)
+export(TXT_NORM, IMG_NORM)

onnx/export/tar.bz2.py

@@ -7,12 +7,12 @@
 
 
 def txz(src, to):
-    stream = bz2.BZ2File(to, 'w')
+  stream = bz2.BZ2File(to, 'w')
 
-    with tarfile.TarFile(fileobj=stream, mode='w') as tar:
-        tar.add(src, arcname=basename(src))
+  with tarfile.TarFile(fileobj=stream, mode='w') as tar:
+    tar.add(src, arcname=basename(src))
 
-    stream.close()
+  stream.close()
 
 
 txz(ONNX_DIR, ONNX_DIR + '.tar.bz2')

onnx/misc/clip_model.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 import torch
+from misc.norm import norm
 import torch.nn as nn
 from .device import DEVICE
 from .config import MODEL_FP
@@ -12,10 +13,10 @@
 MODEL.to(DEVICE)
 
 
-class ImgModel(nn.Module):
+class Img(nn.Module):
 
   def __init__(self):
-    super(ImgModel, self).__init__()
+    super(Img, self).__init__()
     self.model = MODEL
 
   def forward(self, image):
@@ -24,10 +25,16 @@ def forward(self, image):
       return self.model.get_image_features(image)
 
 
-class TxtModel(nn.Module):
+class ImgNorm(Img):
+
+  def forward(self, image):
+    return norm(super(ImgNorm, self).forward(image))
+
+
+class Txt(nn.Module):
 
   def __init__(self):
-    super(TxtModel, self).__init__()
+    super(Txt, self).__init__()
     self.model = MODEL
 
   def forward(self, text, attention_mask):
@@ -37,5 +44,14 @@ def forward(self, text, attention_mask):
       return self.model.get_text_features(text, attention_mask=attention_mask)
 
 
-IMG = ImgModel()
-TXT = TxtModel()
+class TxtNorm(Txt):
+
+  def forward(self, text, attention_mask):
+    return norm(super(TxtNorm, self).forward(text, attention_mask))
+
+
+IMG = Img()
+IMG_NORM = ImgNorm()
+
+TXT = Txt()
+TXT_NORM = TxtNorm()

onnx/misc/config.py

@@ -7,7 +7,7 @@
 
 ROOT = dirname(dirname(abspath(__file__)))
 
-IMG_DIR = join(ROOT, 'img')
+IMG_DIR = join(ROOT, 'test/img')
 MODEL_DIR = join(ROOT, 'model')
 
 MODEL_FP = join(MODEL_DIR, MODEL_NAME)

onnx/misc/norm.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+import torch
+
+
+# 对特征进行归一化
+def norm(vec):
+  with torch.no_grad():
+    vec /= vec.norm(dim=-1, keepdim=True)
+    return vec

onnx/test/clip/clip_img.py

@@ -13,8 +13,12 @@ def img2vec(img):
 
 if __name__ == "__main__":
   from misc.config import IMG_DIR
+  from misc.norm import norm
   from os.path import join
-  fp = join(IMG_DIR, 'cat.jpg')
   from PIL import Image
+
+  fp = join(IMG_DIR, 'cat.jpg')
   img = Image.open(fp)
-  print(img2vec(img))
+  vec = img2vec(img)
+  print('vec', vec)
+  print('norm', norm(vec))

onnx/test/clip/clip_txt.py

@@ -13,11 +13,13 @@ def txt2vec(li):
   from glob import glob
   from misc.config import ROOT
   from test_txt import TEST_TXT
+  from misc.norm import norm
 
   li = glob(join(ROOT, 'jpg/*.jpg'))
   for li in TEST_TXT:
     r = txt2vec(li)
-    for txt, i in zip(li, r):
+    for txt, vec in zip(li, r):
       print(txt)
-      print(i)
+      print('vec', vec)
+      print('norm', norm(vec))
       print('\n')