Merge pull request #336 from modelscope/master_merge_github_0619

Master merge GitHub 0619 from internal

.dev_scripts/build_base_image.sh

@@ -0,0 +1,119 @@
+#!/bin/bash
+# default values.
+BASE_CPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04
+BASE_GPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04-cuda11.3.0-cudnn8-devel
+MODELSCOPE_REPO_ADDRESS=reg.docker.alibaba-inc.com/modelscope/modelscope
+python_version=3.7.13
+torch_version=1.11.0
+cudatoolkit_version=11.3
+tensorflow_version=1.15.5
+version=None
+is_cpu=False
+function usage(){
+    echo "usage: build.sh "
+    echo "       --python=python_version set python version, default: $python_version"
+    echo "       --torch=torch_version set pytorch version, fefault: $torch_version"
+    echo "       --tensorflow=tensorflow_version set tensorflow version, default: $tensorflow_version"
+    echo "       --version=version set image version, default: $version"
+    echo "       --test option for run test before push image, only push on ci test pass"
+    echo "       --cpu option for build cpu version"
+    echo "       --dsw option for build dsw version"
+    echo "       --ci  option for build ci version"
+    echo "       --push option for push image to remote repo"
+}
+for i in "$@"; do
+  case $i in
+    --python=*)
+      python_version="${i#*=}"
+      shift
+      ;;
+    --torch=*)
+      torch_version="${i#*=}"
+      shift # pytorch version
+      ;;
+    --tensorflow=*)
+      tensorflow_version="${i#*=}"
+      shift # tensorflow version
+      ;;
+    --version=*)
+      version="${i#*=}"
+      shift # version
+      ;;
+    --cpu)
+      is_cpu=True
+      shift # is cpu image
+      ;;
+    --push)
+      is_push=True
+      shift # option for push image to remote repo
+      ;;
+    --help)
+      usage
+      exit 0
+      ;;
+    -*|--*)
+      echo "Unknown option $i"
+      usage
+      exit 1
+      ;;
+    *)
+      ;;
+  esac
+done
+
+if [ "$version" == "None" ]; then
+    echo "version must specify!"
+    exit 1
+fi
+if [ "$is_cpu" == "True" ]; then
+    export BASE_IMAGE=$BASE_CPU_IMAGE
+    base_tag=ubuntu20.04
+    export USE_GPU=False
+else
+    export BASE_IMAGE=$BASE_GPU_IMAGE
+    base_tag=ubuntu20.04-cuda11.3.0
+    export USE_GPU=True
+fi
+if [[ $python_version == 3.7* ]]; then
+    base_tag=$base_tag-py37
+elif [[ $python_version == 3.8* ]]; then
+    base_tag=$base_tag-py38
+elif [[ $python_version == 3.9* ]]; then
+    base_tag=$base_tag-py39
+else
+    echo "Unsupport python version: $python_version"
+    exit 1
+fi
+
+target_image_tag=$base_tag-torch$torch_version-tf$tensorflow_version-base-$version
+export IMAGE_TO_BUILD=$MODELSCOPE_REPO_ADDRESS:$target_image_tag
+export PYTHON_VERSION=$python_version
+export TORCH_VERSION=$torch_version
+export CUDATOOLKIT_VERSION=$cudatoolkit_version
+export TENSORFLOW_VERSION=$tensorflow_version
+echo -e "Building image with:\npython$python_version\npytorch$torch_version\ntensorflow:$tensorflow_version\ncudatoolkit:$cudatoolkit_version\ncpu:$is_cpu\n"
+docker_file_content=`cat docker/Dockerfile.ubuntu_base`
+printf "$docker_file_content" > Dockerfile
+
+while true
+do
+  docker build -t $IMAGE_TO_BUILD  \
+             --build-arg USE_GPU \
+             --build-arg BASE_IMAGE \
+             --build-arg PYTHON_VERSION \
+             --build-arg TORCH_VERSION \
+             --build-arg CUDATOOLKIT_VERSION \
+             --build-arg TENSORFLOW_VERSION \
+             -f Dockerfile .
+  if [ $? -eq 0 ]; then
+    echo "Image build done"
+    break
+  else
+    echo "Running docker build command error, we will retry"
+  fi
+done
+
+if [ "$is_push" == "True" ]; then
+    echo "Pushing image: $IMAGE_TO_BUILD"
+    docker push $IMAGE_TO_BUILD
+fi

.dev_scripts/build_image.sh

@@ -1,7 +1,9 @@
 #!/bin/bash
 # default values.
-BASE_CPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04
-BASE_GPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04-cuda11.3.0-cudnn8-devel
+BASE_PY38_CPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/modelscope:ubuntu20.04-py38-torch1.11.0-tf1.15.5-base-1.6.1
+BASE_PY38_GPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/modelscope:ubuntu20.04-cuda11.3.0-py38-torch1.11.0-tf1.15.5-base-1.6.1
+BASE_PY37_CPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/modelscope:ubuntu20.04-py37-torch1.11.0-tf1.15.5-base-1.6.1
+BASE_PY37_GPU_IMAGE=reg.docker.alibaba-inc.com/modelscope/modelscope:ubuntu20.04-cuda11.3.0-py37-torch1.11.0-tf1.15.5-base-1.6.1
 MODELSCOPE_REPO_ADDRESS=reg.docker.alibaba-inc.com/modelscope/modelscope
 python_version=3.7.13
 torch_version=1.11.0
@@ -86,20 +88,30 @@ if [ "$modelscope_version" == "None" ]; then
     exit 1
 fi
 if [ "$is_cpu" == "True" ]; then
-    export BASE_IMAGE=$BASE_CPU_IMAGE
     base_tag=ubuntu20.04
     export USE_GPU=False
 else
-    export BASE_IMAGE=$BASE_GPU_IMAGE
     base_tag=ubuntu20.04-cuda11.3.0
     export USE_GPU=True
 fi
 if [[ $python_version == 3.7* ]]; then
+    if [ "$is_cpu" == "True" ]; then
+        echo "Building python3.7 cpu image"
+        export BASE_IMAGE=$BASE_PY37_CPU_IMAGE
+    else
+        echo "Building python3.7 gpu image"
+        export BASE_IMAGE=$BASE_PY37_GPU_IMAGE
+    fi
     base_tag=$base_tag-py37
 elif [[ $python_version == 3.8* ]]; then
+    if [ "$is_cpu" == "True" ]; then
+        echo "Building python3.8 cpu image"
+        export BASE_IMAGE=$BASE_PY38_CPU_IMAGE
+    else
+        echo "Building python3.8 gpu image"
+        export BASE_IMAGE=$BASE_PY38_GPU_IMAGE
+    fi
     base_tag=$base_tag-py38
-elif [[ $python_version == 3.9* ]]; then
-    base_tag=$base_tag-py39
 else
     echo "Unsupport python version: $python_version"
     exit 1
@@ -120,7 +132,7 @@ echo -e "Building image with:\npython$python_version\npytorch$torch_version\nten
 docker_file_content=`cat docker/Dockerfile.ubuntu`
 if [ "$is_ci_test" != "True" ]; then
     echo "Building ModelScope lib, will install ModelScope lib to image"
-    docker_file_content="${docker_file_content} \nRUN pip install --no-cache-dir  modelscope==$modelscope_version -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html"
+    docker_file_content="${docker_file_content} \nRUN pip install --no-cache-dir modelscope==$modelscope_version -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html"
 fi
 echo "$is_dsw"
 if [ "$is_dsw" == "False" ]; then
@@ -128,6 +140,8 @@ if [ "$is_dsw" == "False" ]; then
 else
     echo "Building dsw image will need set ModelScope lib cache location."
     docker_file_content="${docker_file_content} \nENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope"
+    # pre compile extension
+    docker_file_content="${docker_file_content} \nRUN python -c 'from modelscope.utils.pre_compile import pre_compile_all;pre_compile_all()'"
 fi
 if [ "$is_ci_test" == "True" ]; then
     echo "Building CI image, uninstall modelscope"

.dev_scripts/ci_container_test.sh

@@ -32,6 +32,8 @@ if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then
 else
     echo "Running case in release image, run case directly!"
 fi
+# remove torch_extensions folder to avoid ci hang.
+rm -rf ~/.cache/torch_extensions
 if [ $# -eq 0 ]; then
     ci_command="python tests/run.py --subprocess"
 else

.gitignore

@@ -129,6 +129,5 @@ result.mp4
 *.pth
 *.pt
 
-
 # ast template
 ast_index_file.py

docker/Dockerfile.ubuntu

@@ -1,102 +1,5 @@
-ARG BASE_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04-cuda11.3.0-cudnn8-devel
+ARG BASE_IMAGE=reg.docker.alibaba-inc.com/modelscope/modelscope:ubuntu20.04-cuda11.3.0-py37-torch1.11.0-tf1.15.5-1.6.1
 FROM $BASE_IMAGE
-ARG DEBIAN_FRONTEND=noninteractive
-ENV TZ=Asia/Shanghai
-ENV CONDA_DIR /opt/conda
-ENV PATH="${CONDA_DIR}/bin:${PATH}"
-ENV arch=x86_64
-SHELL ["/bin/bash", "-c"]
-COPY docker/rcfiles /tmp/resources
-COPY docker/jupyter_plugins /tmp/resources/jupyter_plugins
-RUN apt-get update && apt-get install -y --reinstall ca-certificates && \
-    apt-get clean && \
-    cp /tmp/resources/ubuntu20.04_sources.tuna /etc/apt/sources.list && \
-    apt-get update && \
-    apt-get install -y locales wget git strace gdb sox libopenmpi-dev curl strace vim ffmpeg libsm6 tzdata language-pack-zh-hans ttf-wqy-microhei ttf-wqy-zenhei xfonts-wqy libxext6 build-essential ninja-build && \
-    wget https://packagecloud.io/github/git-lfs/packages/debian/bullseye/git-lfs_3.2.0_amd64.deb/download -O ./git-lfs_3.2.0_amd64.deb && \
-    dpkg -i ./git-lfs_3.2.0_amd64.deb && \
-    rm -f ./git-lfs_3.2.0_amd64.deb && \
-    locale-gen zh_CN && \
-    locale-gen zh_CN.utf8 && \
-    update-locale LANG=zh_CN.UTF-8 LC_ALL=zh_CN.UTF-8 LANGUAGE=zh_CN.UTF-8 && \
-    ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
-    dpkg-reconfigure --frontend noninteractive tzdata && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
-
-ENV LANG=zh_CN.UTF-8 LANGUAGE=zh_CN.UTF-8 LC_ALL=zh_CN.UTF-8
-
-#install and config python
-ARG PYTHON_VERSION=3.7.13
-RUN wget --quiet https://mirrors.aliyun.com/anaconda/miniconda/Miniconda3-latest-Linux-${arch}.sh -O ./miniconda.sh && \
-    /bin/bash  miniconda.sh -b -p /opt/conda && \
-    rm  -f miniconda.sh && \
-    ln  -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
-    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
-    cp /tmp/resources/conda.tuna  ~/.condarc && \
-    source /root/.bashrc && \
-    conda install --yes python==${PYTHON_VERSION} && \
-    pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
-    pip config set install.trusted-host mirrors.aliyun.com
-
-ARG USE_GPU=True
-
-# install pytorch
-ARG TORCH_VERSION=1.12.0
-ARG CUDATOOLKIT_VERSION=11.3
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        pip install --no-cache-dir torch==$TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113; \
-    else \
-        pip install --no-cache-dir torch==$TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu; \
-    fi
-
-# install tensorflow
-ARG TENSORFLOW_VERSION=1.15.5
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
-    else \
-        pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION; \
-    fi
-
-# mmcv-full<=1.7.0 for mmdet3d compatible
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="5.0 5.2 6.0 6.1 7.0 7.5 8.0 8.6" MMCV_WITH_OPS=1 MAX_JOBS=8 FORCE_CUDA=1 pip install --no-cache-dir 'mmcv-full<=1.7.0' && pip cache purge; \
-    else \
-        MMCV_WITH_OPS=1 MAX_JOBS=8 pip install --no-cache-dir 'mmcv-full<=1.7.0' && pip cache purge; \
-    fi
-
-# default shell bash
-ENV SHELL=/bin/bash
-# install special package
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        pip install --no-cache-dir dgl-cu113 dglgo -f https://data.dgl.ai/wheels/repo.html; \
-    else \
-        pip install --no-cache-dir dgl dglgo -f https://data.dgl.ai/wheels/repo.html; \
-    fi
-
-# copy install scripts
-COPY docker/scripts/install_unifold.sh docker/scripts/install_colmap.sh docker/scripts/install_pytorch3d_nvdiffrast.sh docker/scripts/install_tiny_cuda_nn.sh docker/scripts/install_apex.sh /tmp/
-
-# for uniford
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        bash /tmp/install_unifold.sh; \
-    else \
-     echo 'cpu unsupport uniford'; \
-    fi
-
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        pip install --no-cache-dir git+https://github.com/gxd1994/Pointnet2.PyTorch.git@master#subdirectory=pointnet2; \
-    else \
-     echo 'cpu unsupport Pointnet2'; \
-    fi
-
-RUN pip install --no-cache-dir detectron2==0.3 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
-
-# 3d supports
-RUN bash /tmp/install_colmap.sh
-RUN bash /tmp/install_tiny_cuda_nn.sh
-RUN bash /tmp/install_pytorch3d_nvdiffrast.sh
-# end of 3D
 
 # install modelscope
 COPY requirements /var/modelscope
@@ -115,12 +18,25 @@ RUN mkdir -p /root/.local/share/jupyter/labextensions/ && \
     cp -r  /tmp/resources/jupyter_plugins/*  /root/.local/share/jupyter/labextensions/
 
 COPY docker/scripts/modelscope_env_init.sh /usr/local/bin/ms_env_init.sh
-RUN pip install --no-cache-dir xtcocotools==1.12  detectron2==0.3 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --force
+# python3.8 pip install git+https://github.com/jin-s13/xtcocoapi.git@v1.13
+# pip install git+https://github.com/gatagat/lap.git@v0.4.0
+RUN pip install --no-cache-dir text2sql_lgesql==1.3.0 \
+         git+https://github.com/jin-s13/xtcocoapi.git@v1.13 \
+         git+https://github.com/gatagat/lap.git@v0.4.0 \
+         detectron2==0.3 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --force --no-deps
 
-# speechbrain==0.5.7 for audio compatible
-RUN pip install --no-cache-dir speechbrain==0.5.7 adaseq>=0.5.0 mmcls>=0.21.0 mmdet>=2.25.0 decord>=0.6.0 numpy==1.18.5 wenetruntime==1.11.0 ipykernel fairseq fasttext deepspeed
+RUN pip install --no-cache-dir mpi4py paint_ldm adaseq>=0.5.0 \
+         mmcls>=0.21.0 mmdet>=2.25.0 decord>=0.6.0 wenetruntime==1.11.0 \
+         ipykernel fairseq fasttext deepspeed -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
+
+# for cpu install cpu version faiss, faiss depends on blas lib, we install libopenblas TODO rename gpu or cpu version faiss
 RUN if [ "$USE_GPU" = "True" ] ; then \
-        bash /tmp/install_apex.sh; \
+        pip install --no-cache-dir funtextprocessing kwsbp==0.0.6 faiss==1.7.2 safetensors typeguard==2.13.3 scikit-learn 'pandas<1.4.0' pai-easycv librosa==0.9.2 funasr -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
     else \
-     echo 'cpu unsupport apex'; \
+        pip install --no-cache-dir funtextprocessing kwsbp==0.0.6 https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/faiss-1.7.2-py37-none-linux_x86_64.whl safetensors typeguard==2.13.3 scikit-learn 'pandas<1.4.0' pai-easycv librosa==0.9.2 funasr -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
     fi
+
+COPY examples /modelscope/examples
+
+# for pai-easycv setup compatiblity issue
+ENV SETUPTOOLS_USE_DISTUTILS=stdlib

docker/Dockerfile.ubuntu_base

@@ -0,0 +1,136 @@
+ARG BASE_IMAGE=reg.docker.alibaba-inc.com/modelscope/ubuntu:20.04-cuda11.3.0-cudnn8-devel
+FROM $BASE_IMAGE
+ARG DEBIAN_FRONTEND=noninteractive
+ENV TZ=Asia/Shanghai
+ENV CONDA_DIR /opt/conda
+ENV PATH="${CONDA_DIR}/bin:${PATH}"
+ENV arch=x86_64
+SHELL ["/bin/bash", "-c"]
+COPY docker/rcfiles /tmp/resources
+COPY docker/jupyter_plugins /tmp/resources/jupyter_plugins
+RUN apt-get update && apt-get install -y --reinstall ca-certificates && \
+    apt-get clean && \
+    cp /tmp/resources/sources.list.aliyun /etc/apt/sources.list && \
+    apt-get update && \
+    apt-get install -y locales wget git strace gdb sox libopenmpi-dev curl \
+    libgeos-dev strace vim ffmpeg libsm6 tzdata language-pack-zh-hans \
+    ttf-wqy-microhei ttf-wqy-zenhei xfonts-wqy libxext6 build-essential ninja-build && \
+    wget https://packagecloud.io/github/git-lfs/packages/debian/bullseye/git-lfs_3.2.0_amd64.deb/download -O ./git-lfs_3.2.0_amd64.deb && \
+    dpkg -i ./git-lfs_3.2.0_amd64.deb && \
+    rm -f ./git-lfs_3.2.0_amd64.deb && \
+    locale-gen zh_CN && \
+    locale-gen zh_CN.utf8 && \
+    update-locale LANG=zh_CN.UTF-8 LC_ALL=zh_CN.UTF-8 LANGUAGE=zh_CN.UTF-8 && \
+    ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
+    dpkg-reconfigure --frontend noninteractive tzdata && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+ENV LANG=zh_CN.UTF-8 LANGUAGE=zh_CN.UTF-8 LC_ALL=zh_CN.UTF-8
+
+#install and config python
+ARG PYTHON_VERSION=3.7.13
+# Miniconda3-py37_23.1.0-1-Linux-x86_64.sh is last python3.7 version
+RUN if [ "$PYTHON_VERSION" = "3.7.13" ] ; then \
+    wget --quiet https://mirrors.aliyun.com/anaconda/miniconda/Miniconda3-py37_23.1.0-1-Linux-x86_64.sh -O ./miniconda.sh && \
+    /bin/bash  miniconda.sh -b -p /opt/conda && \
+    rm  -f miniconda.sh && \
+    ln  -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
+    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
+    cp /tmp/resources/conda.tuna  ~/.condarc && \
+    source /root/.bashrc && \
+    conda install --yes python==${PYTHON_VERSION} && \
+    pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
+    pip config set install.trusted-host mirrors.aliyun.com;\
+else \
+    wget --quiet https://mirrors.aliyun.com/anaconda/miniconda/Miniconda3-latest-Linux-${arch}.sh -O ./miniconda.sh && \
+    /bin/bash  miniconda.sh -b -p /opt/conda && \
+    rm  -f miniconda.sh && \
+    ln  -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
+    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
+    cp /tmp/resources/conda.tuna  ~/.condarc && \
+    source /root/.bashrc && \
+    conda install --yes python==${PYTHON_VERSION} && \
+    pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
+    pip config set install.trusted-host mirrors.aliyun.com;\
+fi
+
+ARG USE_GPU=True
+
+# install pytorch
+ARG TORCH_VERSION=1.12.0
+ARG CUDATOOLKIT_VERSION=cu113
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        pip install --no-cache-dir torch==$TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113; \
+    else \
+        pip install --no-cache-dir torch==$TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu; \
+    fi
+
+# install tensorflow
+ARG TENSORFLOW_VERSION=1.15.5
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html; \
+    else \
+        # only python 3.7 has tensorflow 1.15.5
+        if [ "$PYTHON_VERSION" = "3.7.13" ] ; then \
+            pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION; \
+        else \
+            pip install --no-cache-dir numpy==1.18.5 https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/tensorflow-1.15.5-cp38-cp38-linux_x86_64.whl; \
+         fi \
+    fi
+
+# mmcv-full<=1.7.0 for mmdet3d compatible
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="5.0 5.2 6.0 6.1 7.0 7.5 8.0 8.6" MMCV_WITH_OPS=1 MAX_JOBS=8 FORCE_CUDA=1 pip install --no-cache-dir 'mmcv-full<=1.7.0' && pip cache purge; \
+    else \
+        MMCV_WITH_OPS=1 MAX_JOBS=8 pip install --no-cache-dir 'mmcv-full<=1.7.0' && pip cache purge; \
+    fi
+
+# default shell bash
+ENV SHELL=/bin/bash
+# install special package
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        pip install --no-cache-dir dgl-cu113 dglgo -f https://data.dgl.ai/wheels/repo.html; \
+    else \
+        pip install --no-cache-dir dgl==0.9.0 dglgo -f https://data.dgl.ai/wheels/repo.html; \
+    fi
+
+# copy install scripts
+COPY docker/scripts/install_unifold.sh docker/scripts/install_colmap.sh docker/scripts/install_pytorch3d_nvdiffrast.sh docker/scripts/install_tiny_cuda_nn.sh docker/scripts/install_apex.sh /tmp/
+
+# for uniford
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        bash /tmp/install_unifold.sh; \
+    else \
+     echo 'cpu unsupport uniford'; \
+    fi
+
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        pip install --no-cache-dir git+https://github.com/gxd1994/Pointnet2.PyTorch.git@master#subdirectory=pointnet2; \
+    else \
+     echo 'cpu unsupport Pointnet2'; \
+    fi
+
+# 3d supports
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        bash /tmp/install_colmap.sh; \
+    else \
+     echo 'cpu unsupport colmap'; \
+    fi
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        bash /tmp/install_tiny_cuda_nn.sh \
+    else \
+     echo 'cpu unsupport tiny_cudann'; \
+    fi
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        bash /tmp/install_pytorch3d_nvdiffrast.sh; \
+    else \
+     echo 'cpu unsupport pytorch3d nvdiffrast'; \
+    fi
+# end of 3D
+# install apex after deepspeed
+RUN if [ "$USE_GPU" = "True" ] ; then \
+        bash /tmp/install_apex.sh; \
+    else \
+     echo 'cpu unsupport apex'; \
+    fi

docker/rcfiles/sources.list.aliyun

@@ -1,25 +1,14 @@
-deb http://mirrors.aliyun.com/ubuntu/ bionic main restricted
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic main restricted
+deb https://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
+# deb-src https://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
 
-deb http://mirrors.aliyun.com/ubuntu/ bionic-updates main restricted
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic-updates main restricted
+deb https://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
+# deb-src https://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
 
-deb http://mirrors.aliyun.com/ubuntu/ bionic universe
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic universe
-deb http://mirrors.aliyun.com/ubuntu/ bionic-updates universe
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic-updates universe
+deb https://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
+# deb-src https://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
 
-deb http://mirrors.aliyun.com/ubuntu/ bionic multiverse
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic multiverse
-deb http://mirrors.aliyun.com/ubuntu/ bionic-updates multiverse
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic-updates multiverse
+# deb https://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
+# deb-src https://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
 
-deb http://mirrors.aliyun.com/ubuntu/ bionic-backports main restricted universe multiverse
-# deb-src http://mirrors.aliyun.com/ubuntu/ bionic-backports main restricted universe multiverse
-
-deb http://mirrors.aliyun.com/ubuntu bionic-security main restricted
-# deb-src http://mirrors.aliyun.com/ubuntu bionic-security main restricted
-deb http://mirrors.aliyun.com/ubuntu bionic-security universe
-# deb-src http://mirrors.aliyun.com/ubuntu bionic-security universe
-deb http://mirrors.aliyun.com/ubuntu bionic-security multiverse
-# deb-src http://mirrors.aliyun.com/ubuntu bionic-security multiverse
+deb https://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse
+# deb-src https://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse

docker/scripts/install_apex.sh

@@ -1,6 +1,7 @@
 export MAX_JOBS=16 \
 && git clone https://github.com/NVIDIA/apex \
 && cd apex \
-&& TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.0;8.6" pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ \
+&& git checkout 6bd01c4b99a84648ad5e5238a959735e6936c813 \
+&& TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.0;8.6" pip install -v --disable-pip-version-check --no-cache --global-option="--cpp_ext" --global-option="--cuda_ext" ./ \
 && cd .. \
 && rm -rf apex

docs/source/api/modelscope.pipelines.multi_modal.rst

@@ -18,7 +18,7 @@ modelscope.pipelines.multi_modal
     ImageCaptioningPipeline
     MGeoRankingPipeline
     MultiModalEmbeddingPipeline
-    StableDiffusionWrapperPipeline
+    StableDiffusionPipeline
     TextToImageSynthesisPipeline
     VideoCaptioningPipeline
     VideoMultiModalEmbeddingPipeline

examples/pytorch/llama/finetune_llama.py

@@ -0,0 +1,263 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import copy
+import logging
+import os
+import shutil
+import tempfile
+import unittest
+from dataclasses import dataclass, field
+
+import json
+import torch
+import utils
+
+from modelscope import TrainingArgs
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.metainfo import Trainers
+from modelscope.models.nlp.llama import LlamaForTextGeneration, LlamaTokenizer
+from modelscope.msdatasets.dataset_cls.custom_datasets.torch_custom_dataset import \
+    TorchCustomDataset
+from modelscope.trainers import build_trainer
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = '[PAD]'
+DEFAULT_EOS_TOKEN = '</s>'
+DEFAULT_BOS_TOKEN = '<s>'
+DEFAULT_UNK_TOKEN = '<unk>'
+PROMPT_DICT = {
+    'prompt_input':
+    ('Below is an instruction that describes a task, paired with an input that provides further context. '
+     'Write a response that appropriately completes the request.\n\n'
+     '### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:'
+     ),
+    'prompt_no_input':
+    ('Below is an instruction that describes a task. '
+     'Write a response that appropriately completes the request.\n\n'
+     '### Instruction:\n{instruction}\n\n### Response:'),
+}
+
+
+@dataclass(init=False)
+class TextGenerationArguments(TrainingArgs):
+    src_txt: str = field(
+        default=None,
+        metadata={
+            'help': 'The source text key of preprocessor',
+            'cfg_node': 'preprocessor.src_txt'
+        })
+
+    deepspeed: str = field(
+        default=None,
+        metadata={
+            'help': 'The location of DeepSpeed json config file.',
+        })
+
+    work_dir: str = field(
+        default=None, metadata={
+            'help': 'The location of work dir',
+        })
+
+
+def _tokenize_fn(strings, tokenizer):
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors='pt',
+            padding='longest',
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        ) for text in strings
+    ]
+    input_ids = labels = [
+        tokenized.input_ids[0] for tokenized in tokenized_list
+    ]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()
+        for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+def preprocess(sources, targets, tokenizer):
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [
+        _tokenize_fn(strings, tokenizer) for strings in (examples, sources)
+    ]
+    input_ids = examples_tokenized['input_ids']
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized['input_ids_lens']):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+
+def smart_tokenizer_and_embedding_resize(special_tokens_dict, tokenizer,
+                                         model):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+class SupervisedDataset(TorchCustomDataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(self, data_path: str, tokenizer):
+        logging.warning('Loading data...')
+        f = open(data_path, 'r')
+        list_data_dict = json.load(f)
+        f.close()
+
+        logging.warning('Formatting inputs...')
+        prompt_input, prompt_no_input = PROMPT_DICT[
+            'prompt_input'], PROMPT_DICT['prompt_no_input']
+        sources = [
+            prompt_input.format_map(example) if example.get('input', '') != ''
+            else prompt_no_input.format_map(example)
+            for example in list_data_dict
+        ]
+        targets = [
+            f"{example['output']}{tokenizer.eos_token}"
+            for example in list_data_dict
+        ]
+
+        logging.warning('Tokenizing inputs... This may take some time...')
+        data_dict = preprocess(sources, targets, tokenizer)
+
+        self.input_ids = data_dict['input_ids']
+        self.labels = data_dict['labels']
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, i):
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: LlamaTokenizer
+
+    def __call__(self, instances):
+        input_ids, labels = tuple([instance[key] for instance in instances]
+                                  for key in ('input_ids', 'labels'))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids,
+            batch_first=True,
+            padding_value=self.tokenizer.pad_token_id)
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+
+config, args = TextGenerationArguments().parse_cli().to_config()
+
+if __name__ == '__main__':
+
+    def cfg_modify_fn(cfg):
+        if args.use_model_config:
+            cfg.merge_from_dict(config)
+        else:
+            cfg = config
+        cfg.train.lr_scheduler = {
+            'type': 'CosineAnnealingLR',
+            'T_max': 1,
+            'options': {
+                'by_epoch': False
+            }
+        }
+        cfg.train.optimizer = {
+            'type': 'AdamW',
+            'lr': 2e-5,
+            'weight_decay': 0.0,
+            'options': {
+                'cumulative_iters': 8,
+                'warmup': {
+                    'type': 'LinearWarmup',
+                    'warmup_ratio': 0.03
+                }
+            }
+        }
+        cfg.train.logging = {'interval': 8, 'by_epoch': False}
+        cfg.train['bf16'] = True
+        cfg.train.dataloader = {'batch_size_per_gpu': 4, 'workers_per_gpu': 1}
+        if 'hooks' not in cfg.train:
+            cfg.train['hooks'] = []
+        cfg.train.hooks.append({
+            'type': 'DeepspeedHook',
+            'config': args.deepspeed,
+            'save_zero_checkpoint': True,
+            'with_mpu': False,
+        })
+
+        cfg.preprocessor.sequence_length = 512
+        return cfg
+
+    model_path = args.model if os.path.exists(
+        args.model) else snapshot_download(args.model)
+    data_path = args.src_txt if args.src_txt else os.path.join(
+        model_path, 'alpaca_data.json')
+    model = LlamaForTextGeneration.from_pretrained(model_path)
+
+    tokenizer = LlamaTokenizer.from_pretrained(
+        model_path,
+        model_max_length=512,
+        padding_side='right',
+    )
+
+    special_tokens_dict = dict()
+    special_tokens_dict['pad_token'] = DEFAULT_PAD_TOKEN
+    special_tokens_dict['eos_token'] = DEFAULT_EOS_TOKEN
+    special_tokens_dict['bos_token'] = DEFAULT_BOS_TOKEN
+    special_tokens_dict['unk_token'] = DEFAULT_UNK_TOKEN
+
+    smart_tokenizer_and_embedding_resize(
+        special_tokens_dict=special_tokens_dict,
+        tokenizer=tokenizer,
+        model=model,
+    )
+
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=data_path)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+
+    kwargs = dict(
+        model=model,
+        cfg_file=os.path.join(model_path, 'configuration.json'),
+        train_dataset=train_dataset,
+        data_collator=data_collator,
+        max_epochs=3,
+        work_dir=args.work_dir,
+        cfg_modify_fn=cfg_modify_fn)
+
+    # Construct trainer and train
+    trainer = build_trainer(
+        name=Trainers.text_generation_trainer, default_args=kwargs)
+    trainer.train()

examples/pytorch/llama/run_train_llama.sh

@@ -0,0 +1,9 @@
+DATA_PARALLEL_SIZE=4
+
+
+export PYTHONPATH=$PYTHONPATH:./
+torchrun --nproc_per_node $DATA_PARALLEL_SIZE examples/pytorch/llama/finetune_llama.py \
+    --work_dir './tmp' \
+    --model 'skyline2006/llama-7b' \
+    --deepspeed 'default_offload_opt_param.json' \
+    --eval_interval 100

examples/pytorch/named_entity_recognition/finetune_named_entity_recognition.py

@@ -0,0 +1,97 @@
+import os
+from dataclasses import dataclass, field
+
+from adaseq.data.data_collators.base import build_data_collator
+from adaseq.data.dataset_manager import DatasetManager
+from adaseq.data.preprocessors.nlp_preprocessor import build_preprocessor
+from adaseq.training.default_trainer import DefaultTrainer as AdaSeqTrainer
+
+from modelscope import MsDataset, TrainingArgs, build_dataset_from_file
+
+
+@dataclass(init=False)
+class NamedEntityRecognitionArguments(TrainingArgs):
+    preprocessor: str = field(
+        default='sequence-labeling-preprocessor',
+        metadata={
+            'help': 'The preprocessor type',
+            'cfg_node': 'preprocessor.type'
+        })
+
+    sequence_length: int = field(
+        default=150,
+        metadata={
+            'cfg_node': 'preprocessor.max_length',
+            'help': 'The parameters for train dataset',
+        })
+
+    data_collator: str = field(
+        default='SequenceLabelingDataCollatorWithPadding',
+        metadata={
+            'cfg_node': 'data_collator',
+            'help': 'The type of data collator',
+        })
+
+    dropout: float = field(
+        default=0.0,
+        metadata={
+            'cfg_node': 'model.dropout',
+            'help': 'Dropout rate',
+        })
+
+    use_crf: bool = field(
+        default=True,
+        metadata={
+            'cfg_node': 'model.use_crf',
+            'help': 'Whether to add a CRF decoder layer',
+        })
+
+    crf_lr: float = field(
+        default=5.0e-1, metadata={
+            'help': 'Learning rate for CRF layer',
+        })
+
+
+training_args = NamedEntityRecognitionArguments().parse_cli()
+config, args = training_args.to_config()
+print(args)
+
+if args.dataset_json_file is None:
+    train_dataset = MsDataset.load(
+        args.train_dataset_name,
+        subset_name=args.train_subset_name,
+        split=args.train_split,
+        namespace=args.train_dataset_namespace).to_hf_dataset()
+    validation_dataset = MsDataset.load(
+        args.val_dataset_name,
+        subset_name=args.val_subset_name,
+        split=args.val_split,
+        namespace=args.val_dataset_namespace).to_hf_dataset()
+else:
+    train_dataset, validation_dataset = build_dataset_from_file(
+        args.dataset_json_file)
+dm = DatasetManager({
+    'train': train_dataset,
+    'valid': validation_dataset
+}, labels={'type': 'count_span_labels'})  # yapf: disable
+
+config.preprocessor.model_dir = args.model
+config.model.embedder = {'model_name_or_path': args.model}
+preprocessor = build_preprocessor(config.preprocessor, labels=dm.labels)
+config.model.id_to_label = preprocessor.id_to_label
+data_collator = build_data_collator(preprocessor.tokenizer,
+                                    dict(type=config.data_collator))
+config.train.optimizer.param_groups = [{'regex': 'crf', 'lr': args.crf_lr}]
+
+cfg_file = os.path.join(config.train.work_dir, 'config.yaml')
+config.dump(cfg_file)
+
+kwargs = dict(
+    cfg_file=cfg_file,
+    work_dir=config.train.work_dir,
+    dataset_manager=dm,
+    data_collator=data_collator,
+    preprocessor=preprocessor)
+
+trainer = AdaSeqTrainer(**kwargs)
+trainer.train()

examples/pytorch/named_entity_recognition/run_train.sh

@@ -0,0 +1,26 @@
+PYTHONPATH=. python examples/pytorch/named_entity_recognition/finetune_named_entity_recognition.py \
+    --task 'named-entity-recognition' \
+    --work_dir './tmp' \
+    --model_type 'sequence-labeling-model' \
+    --model 'damo/nlp_structbert_backbone_base_std' \
+    --dropout 0.1 \
+    --use_crf true \
+    --train_dataset_name 'resume_ner' \
+    --train_dataset_namespace 'damo' \
+    --train_split 'train' \
+    --val_dataset_name 'resume_ner' \
+    --val_dataset_namespace 'damo' \
+    --val_split 'dev' \
+    --preprocessor 'sequence-labeling-preprocessor' \
+    --sequence_length 150 \
+    --data_collator 'SequenceLabelingDataCollatorWithPadding' \
+    --max_epochs 5 \
+    --per_device_train_batch_size 16 \
+    --train_data_worker 0 \
+    --eval_data_worker 0 \
+    --lr 5.0e-5 \
+    --lr_scheduler LinearLR \
+    --lr_scheduler_params 'start_factor=1.0,end_factor=0.0,total_iters=5' \
+    --eval_metrics ner-metric \
+    --save_best_checkpoint true \
+    --metric_for_best_model f1 \

examples/pytorch/stable_diffusion/finetune_stable_diffusion.py

@@ -1,17 +1,21 @@
-from dataclasses import dataclass, field
-
+from modelscope.metainfo import Trainers
 from modelscope.msdatasets import MsDataset
 from modelscope.trainers import EpochBasedTrainer, build_trainer
 from modelscope.trainers.training_args import TrainingArgs
+from modelscope.utils.constant import DownloadMode
 
-training_args = TrainingArgs(task='efficient-diffusion-tuning').parse_cli()
+training_args = TrainingArgs(task='text-to-image-synthesis').parse_cli()
 config, args = training_args.to_config()
 print(args)
 
-dataset = MsDataset.load(
-    args.train_dataset_name, namespace=args.train_dataset_namespace)
-train_dataset = dataset['train']
-validation_dataset = dataset['validation']
+train_dataset = MsDataset.load(
+    args.train_dataset_name,
+    split='train',
+    download_mode=DownloadMode.FORCE_REDOWNLOAD)
+validation_dataset = MsDataset.load(
+    args.train_dataset_name,
+    split='validation',
+    download_mode=DownloadMode.FORCE_REDOWNLOAD)
 
 
 def cfg_modify_fn(cfg):
@@ -19,17 +23,22 @@ def cfg_modify_fn(cfg):
         cfg.merge_from_dict(config)
     else:
         cfg = config
-    cfg.train.lr_scheduler.T_max = training_args.max_epochs
-    cfg.model.inference = False
+    cfg.train.lr_scheduler = {
+        'type': 'LambdaLR',
+        'lr_lambda': lambda _: 1,
+        'last_epoch': -1
+    }
+    cfg.train.optimizer.lr = 1e-4
     return cfg
 
 
 kwargs = dict(
     model=training_args.model,
+    model_revision='v1.0.6',
     work_dir=training_args.work_dir,
     train_dataset=train_dataset,
     eval_dataset=validation_dataset,
     cfg_modify_fn=cfg_modify_fn)
 
-trainer: EpochBasedTrainer = build_trainer(name='trainer', default_args=kwargs)
+trainer = build_trainer(name=Trainers.lora_diffusion, default_args=kwargs)
 trainer.train()

examples/pytorch/stable_diffusion/run_train.sh

@@ -1,12 +1,12 @@
 PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/finetune_stable_diffusion.py \
-    --model 'damo/multi-modal_efficient-diffusion-tuning-lora' \
-    --work_dir './tmp/stable_diffusion_tuning' \
-    --train_dataset_namespace 'damo' \
-    --train_dataset_name 'controlnet_dataset_condition_fill50k' \
-    --max_epochs 1 \
+    --model 'AI-ModelScope/stable-diffusion-v1-5' \
+    --model_revision 'v1.0.6' \
+    --work_dir './tmp/lora_diffusion' \
+    --train_dataset_name 'buptwq/lora-stable-diffusion-finetune' \
+    --max_epochs 100 \
     --save_ckpt_strategy 'by_epoch' \
     --logging_interval 100 \
     --train.dataloader.workers_per_gpu 0 \
     --evaluation.dataloader.workers_per_gpu 0 \
-    --train.optimizer.lr 1e-5 \
+    --train.optimizer.lr 1e-4 \
     --use_model_config true

examples/pytorch/text_classification/finetune_text_classification.py

@@ -70,16 +70,23 @@ def cfg_modify_fn(cfg):
 
 
 if args.dataset_json_file is None:
-    dataset = MsDataset.load(
-        args.train_dataset_name, subset_name=args.train_subset_name)
-    train_dataset = dataset['train']
-    validation_dataset = dataset['validation']
+    train_dataset = MsDataset.load(
+        args.train_dataset_name,
+        subset_name=args.train_subset_name,
+        split=args.train_split,
+        namespace=args.train_dataset_namespace)
+    validation_dataset = MsDataset.load(
+        args.val_dataset_name,
+        subset_name=args.val_subset_name,
+        split=args.val_split,
+        namespace=args.val_dataset_namespace)
 else:
     train_dataset, validation_dataset = build_dataset_from_file(
         args.dataset_json_file)
 
 kwargs = dict(
     model=args.model,
+    model_revision=args.model_revision,
     train_dataset=train_dataset,
     eval_dataset=validation_dataset,
     seed=args.seed,

examples/pytorch/text_classification/run_train.sh

@@ -2,15 +2,24 @@ PYTHONPATH=. python examples/pytorch/text_classification/finetune_text_classific
     --task 'text-classification' \
     --model 'damo/nlp_structbert_backbone_base_std' \
     --train_dataset_name 'clue' \
+    --val_dataset_name 'clue' \
     --train_subset_name 'tnews' \
+    --val_subset_name 'tnews' \
+    --train_split 'train' \
+    --val_split 'validation' \
     --first_sequence 'sentence' \
-    --preprocessor.label label \
-    --model.num_labels 15 \
+    --label label \
+    --num_labels 15 \
     --labels '0,1,2,3,4,5,6,7,8,9,10,11,12,13,14' \
     --preprocessor 'sen-cls-tokenizer' \
     --use_model_config True \
     --max_epochs 1 \
-    --train.dataloader.workers_per_gpu 0 \
-    --evaluation.dataloader.workers_per_gpu 0 \
-    --train.optimizer.lr 1e-5 \
+    --per_device_train_batch_size 16 \
+    --per_device_eval_batch_size 16 \
+    --eval_interval 100 \
+    --eval_strategy by_step \
+    --work_dir './tmp' \
+    --train_data_worker 0 \
+    --eval_data_worker 0 \
+    --lr 1e-5 \
     --eval_metrics 'seq-cls-metric' \

examples/pytorch/text_generation/finetune_text_generation.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 
-from modelscope import EpochBasedTrainer, MsDataset, TrainingArgs
+from modelscope import (EpochBasedTrainer, MsDataset, TrainingArgs,
+                        build_dataset_from_file)
 from modelscope.metainfo import Trainers
 from modelscope.trainers import build_trainer
 
@@ -94,14 +95,26 @@ def cfg_modify_fn(cfg):
     return cfg
 
 
-dataset = MsDataset.load(args.train_dataset_name)
-train_dataset = dataset['train']
-eval_dataset = dataset['validation' if 'validation' in dataset else 'test']
+if args.dataset_json_file is None:
+    train_dataset = MsDataset.load(
+        args.train_dataset_name,
+        subset_name=args.train_subset_name,
+        split=args.train_split,
+        namespace=args.train_dataset_namespace)
+    validation_dataset = MsDataset.load(
+        args.val_dataset_name,
+        subset_name=args.val_subset_name,
+        split=args.val_split,
+        namespace=args.val_dataset_namespace)
+else:
+    train_dataset, validation_dataset = build_dataset_from_file(
+        args.dataset_json_file)
 
 kwargs = dict(
     model=args.model,
+    model_revision=args.model_revision,
     train_dataset=train_dataset,
-    eval_dataset=eval_dataset,
+    eval_dataset=validation_dataset,
     seed=args.seed,
     work_dir=args.work_dir,
     cfg_modify_fn=cfg_modify_fn)

examples/pytorch/text_generation/run_train_gpt3.sh

@@ -9,6 +9,9 @@ PYTHONPATH=. torchrun --nproc_per_node $WORLD_SIZE examples/pytorch/text_generat
     --work_dir './tmp' \
     --model 'damo/nlp_gpt3_text-generation_1.3B' \
     --train_dataset_name 'chinese-poetry-collection' \
+    --val_dataset_name 'chinese-poetry-collection' \
+    --train_split 'train' \
+    --val_split 'test' \
     --preprocessor 'text-gen-jieba-tokenizer' \
     --src_txt 'text1' \
     --tgt_txt 'text2' \

examples/pytorch/text_generation/run_train_mt5.sh

@@ -4,6 +4,9 @@ PYTHONPATH=. torchrun examples/pytorch/text_generation/finetune_text_generation.
     --task 'text2text-generation' \
     --model 'damo/nlp_mt5_zero-shot-augment_chinese-base' \
     --train_dataset_name 'DuReader_robust-QG' \
+    --val_dataset_name 'DuReader_robust-QG' \
+    --train_split 'train' \
+    --val_split 'validation' \
     --src_txt 'text1' \
     --tgt_txt 'text2' \
     --max_epochs 1 \

examples/pytorch/text_generation/run_train_palm.sh

@@ -3,6 +3,9 @@ PYTHONPATH=. torchrun examples/pytorch/text_generation/finetune_text_generation.
     --work_dir './tmp' \
     --model 'damo/nlp_palm2.0_pretrained_chinese-base' \
     --train_dataset_name 'DuReader_robust-QG' \
+    --val_dataset_name 'DuReader_robust-QG' \
+    --train_split 'train' \
+    --val_split 'validation' \
     --src_txt 'text1' \
     --tgt_txt 'text2' \
     --max_epochs 1 \

modelscope/exporters/cv/object_detection_damoyolo_exporter.py

@@ -13,6 +13,9 @@ from modelscope.metainfo import Models
 from modelscope.utils.constant import ModelFile, Tasks
 
 
+@EXPORTERS.register_module(
+    Tasks.domain_specific_object_detection,
+    module_name=Models.tinynas_damoyolo)
 @EXPORTERS.register_module(
     Tasks.image_object_detection, module_name=Models.tinynas_damoyolo)
 class ObjectDetectionDamoyoloExporter(TorchModelExporter):

modelscope/hub/api.py

@@ -16,6 +16,7 @@ from http.cookiejar import CookieJar
 from os.path import expanduser
 from typing import Dict, List, Optional, Tuple, Union
 
+import pandas as pd
 import requests
 from requests import Session
 from requests.adapters import HTTPAdapter, Retry
@@ -46,7 +47,8 @@ from modelscope.utils.constant import (DEFAULT_DATASET_REVISION,
                                        MASTER_MODEL_BRANCH, DatasetFormations,
                                        DatasetMetaFormats,
                                        DatasetVisibilityMap, DownloadChannel,
-                                       ModelFile, VirgoDatasetConfig)
+                                       DownloadMode, ModelFile,
+                                       VirgoDatasetConfig)
 from modelscope.utils.logger import get_logger
 from .utils.utils import (get_endpoint, get_release_datetime,
                           model_id_to_group_owner_name)
@@ -640,16 +642,45 @@ class HubApi:
 
         return local_paths, dataset_formation
 
-    def fetch_single_csv_script(self, script_url: str):
+    @staticmethod
+    def fetch_csv_from_url(url, out_path, chunk_size=100000, mode=DownloadMode.REUSE_DATASET_IF_EXISTS):
+        from io import StringIO
+        import hashlib
+        out_path = os.path.join(out_path, hashlib.md5(url.encode(encoding='UTF-8')).hexdigest())
+        if mode == DownloadMode.FORCE_REDOWNLOAD and os.path.exists(out_path):
+            os.remove(out_path)
+        if os.path.exists(out_path):
+            logger.info(f'Reusing cached meta-csv file: {out_path}')
+            return out_path
         cookies = ModelScopeConfig.get_cookies()
-        resp = self.session.get(script_url, cookies=cookies, headers=self.headers)
-        if not resp or not resp.text:
-            raise 'The meta-csv file cannot be empty when the meta-args `big_data` is true.'
-        text_list = resp.text.strip().split('\n')
-        text_headers = text_list[0]
-        text_content = text_list[1:]
 
-        return text_headers, text_content
+        # Make the request and get the response content as TextIO
+        logger.info('Loading meta-csv file ...')
+
+        response = requests.get(url, cookies=cookies)
+        data = StringIO(response.text)
+
+        # Use read_csv with the TextIO object
+        csv_file_reader = pd.read_csv(data, iterator=True, dtype=str, delimiter=None)
+
+        loop = True
+        iter_num = 0
+        while loop:
+            try:
+                chunk = csv_file_reader.get_chunk(size=chunk_size)
+                logger.info(f'Receiving chunk {iter_num}, shape: {chunk.shape}')
+                if iter_num == 0:
+                    with_header = True
+                else:
+                    with_header = False
+
+                chunk.to_csv(out_path, mode='a', index=False, header=with_header)
+                iter_num += 1
+            except StopIteration:
+                loop = False
+                logger.info('stop chunk iteration')
+
+        return out_path
 
     def get_dataset_file_url(
             self,

modelscope/hub/push_to_hub.py

@@ -2,6 +2,8 @@
 
 import concurrent.futures
 import os
+import shutil
+from multiprocessing import Manager, Process, Value
 
 from modelscope.hub.api import HubApi
 from modelscope.hub.constants import ModelVisibility
@@ -11,6 +13,10 @@ from modelscope.utils.logger import get_logger
 logger = get_logger()
 
 _executor = concurrent.futures.ProcessPoolExecutor(max_workers=8)
+_queues = dict()
+_flags = dict()
+_tasks = dict()
+_manager = None
 
 
 def _api_push_to_hub(repo_name,
@@ -131,3 +137,64 @@ def push_to_hub_async(repo_name,
     return _executor.submit(_api_push_to_hub, repo_name, output_dir, token,
                             private, commit_message, tag, source_repo,
                             ignore_file_pattern, revision)
+
+
+def submit_task(q, b):
+    while True:
+        b.value = False
+        item = q.get()
+        logger.info(item)
+        b.value = True
+        if not item.pop('done', False):
+            delete_dir = item.pop('delete_dir', False)
+            output_dir = item.get('output_dir')
+            try:
+                push_to_hub(**item)
+                if delete_dir and os.path.exists(output_dir):
+                    shutil.rmtree(output_dir)
+            except Exception as e:
+                logger.error(e)
+        else:
+            break
+
+
+class UploadStrategy:
+    cancel = 'cancel'
+    wait = 'wait'
+
+
+def push_to_hub_in_queue(queue_name, strategy=UploadStrategy.cancel, **kwargs):
+    assert queue_name is not None and len(
+        queue_name) > 0, 'Please specify a valid queue name!'
+    global _manager
+    if _manager is None:
+        _manager = Manager()
+    if queue_name not in _queues:
+        _queues[queue_name] = _manager.Queue()
+        _flags[queue_name] = Value('b', False)
+        process = Process(
+            target=submit_task, args=(_queues[queue_name], _flags[queue_name]))
+        process.start()
+        _tasks[queue_name] = process
+
+    queue = _queues[queue_name]
+    flag: Value = _flags[queue_name]
+    if kwargs.get('done', False):
+        queue.put(kwargs)
+    elif flag.value and strategy == UploadStrategy.cancel:
+        logger.error(
+            f'Another uploading is running, '
+            f'this uploading with message {kwargs.get("commit_message")} will be canceled.'
+        )
+    else:
+        queue.put(kwargs)
+
+
+def wait_for_done(queue_name):
+    process: Process = _tasks.pop(queue_name, None)
+    if process is None:
+        return
+    process.join()
+
+    _queues.pop(queue_name)
+    _flags.pop(queue_name)

modelscope/metainfo.py

@@ -205,6 +205,7 @@ class Models(object):
     efficient_diffusion_tuning = 'efficient-diffusion-tuning'
     mplug_owl = 'mplug-owl'
     clip_interrogator = 'clip-interrogator'
+    stable_diffusion = 'stable-diffusion'
 
     # science models
     unifold = 'unifold'
@@ -892,6 +893,8 @@ class MultiModalTrainers(object):
     mplug = 'mplug'
     mgeo_ranking_trainer = 'mgeo-ranking-trainer'
     efficient_diffusion_tuning = 'efficient-diffusion-tuning'
+    stable_diffusion = 'stable-diffusion'
+    lora_diffusion = 'lora-diffusion'
 
 
 class AudioTrainers(object):

modelscope/models/cv/image_human_parsing/parsing_utils.py

@@ -18,7 +18,7 @@ def center_to_target_size_test(img, target_size):
     new_h, new_w = 0, 0
     tfm_list = []
     if src_h > trg_h and src_w > trg_w:
-        if src_h > src_w:
+        if src_h >= src_w:
             new_h = trg_h
             new_w = int(new_h * src_w / src_h)
             if new_w > trg_w:

modelscope/models/cv/ocr_detection/model.py

@@ -12,7 +12,7 @@ from modelscope.models.builder import MODELS
 from modelscope.utils.config import Config
 from modelscope.utils.constant import ModelFile, Tasks
 from modelscope.utils.logger import get_logger
-from .modules.dbnet import DBModel, VLPTModel
+from .modules.dbnet import DBModel, DBNasModel, VLPTModel
 from .utils import boxes_from_bitmap, polygons_from_bitmap
 
 LOGGER = get_logger()
@@ -40,6 +40,8 @@ class OCRDetection(TorchModel):
             self.detector = VLPTModel()
         elif self.backbone == 'resnet18':
             self.detector = DBModel()
+        elif self.backbone == 'proxylessnas':
+            self.detector = DBNasModel()
         else:
             raise TypeError(
                 f'detector backbone should be either resnet18, resnet50, but got {cfgs.model.backbone}'

modelscope/models/cv/ocr_detection/modules/dbnet.py

@@ -10,6 +10,8 @@ from collections import OrderedDict
 import torch
 import torch.nn as nn
 
+from .proxyless import CompactDetBackbone
+
 BatchNorm2d = nn.BatchNorm2d
 
 
@@ -30,6 +32,73 @@ def conv3x3(in_planes, out_planes, stride=1):
         bias=False)
 
 
+class DwPwConv(nn.Module):
+
+    def __init__(self,
+                 in_planes,
+                 out_planes,
+                 kernel_size,
+                 stride=1,
+                 padding=1,
+                 bias=False):
+        super(DwPwConv, self).__init__()
+        self.depthwise = nn.Conv2d(
+            in_planes,
+            in_planes,
+            kernel_size,
+            stride,
+            padding,
+            groups=in_planes,
+            bias=bias)
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.pointwise = nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            groups=1,
+            bias=bias)
+
+    def forward(self, x):
+        out = self.depthwise(x)
+        out = self.bn1(out)
+        out = self.relu1(out)
+
+        out = self.pointwise(out)
+
+        return out
+
+
+class DwPwConvTranspose(nn.Module):
+
+    def __init__(self, in_planes, out_planes, kernel_size, stride):
+        super(DwPwConvTranspose, self).__init__()
+        self.depthwise = nn.ConvTranspose2d(
+            in_planes, in_planes, kernel_size, stride, groups=in_planes)
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.pointwise = nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            groups=1)
+
+    def forward(self, x):
+        out = self.depthwise(x)
+        out = self.bn1(out)
+        out = self.relu1(out)
+
+        out = self.pointwise(out)
+
+        return out
+
+
 class BasicBlock(nn.Module):
     expansion = 1
 
@@ -266,6 +335,156 @@ class ResNet(nn.Module):
         return x2, x3, x4, x5
 
 
+class LightSegDetector(nn.Module):
+
+    def __init__(self,
+                 in_channels=[64, 128, 256, 512],
+                 inner_channels=256,
+                 k=10,
+                 bias=False,
+                 adaptive=False,
+                 smooth=False,
+                 serial=False,
+                 dw_kernel_size=3,
+                 dw_padding=1,
+                 *args,
+                 **kwargs):
+        '''
+        bias: Whether conv layers have bias or not.
+        adaptive: Whether to use adaptive threshold training or not.
+        smooth: If true, use bilinear instead of deconv.
+        serial: If true, thresh prediction will combine segmentation result as input.
+        '''
+        super(LightSegDetector, self).__init__()
+        self.k = k
+        self.serial = serial
+        self.inner_channels = inner_channels
+        self.bias = bias
+        self.dw_kernel_size = dw_kernel_size
+        self.dw_padding = dw_padding
+
+        self.up5 = nn.Upsample(scale_factor=8, mode='nearest')
+        self.up4 = nn.Upsample(scale_factor=4, mode='nearest')
+        self.up3 = nn.Upsample(scale_factor=2, mode='nearest')
+
+        self.in5 = nn.Conv2d(in_channels[-1], inner_channels, 1, bias=bias)
+        self.in4 = nn.Conv2d(in_channels[-2], inner_channels, 1, bias=bias)
+        self.in3 = nn.Conv2d(in_channels[-3], inner_channels, 1, bias=bias)
+        self.in2 = nn.Conv2d(in_channels[-4], inner_channels, 1, bias=bias)
+
+        # DwPM
+        self.binarize = nn.Sequential(
+            DwPwConv(
+                inner_channels,
+                inner_channels // 4,
+                kernel_size=self.dw_kernel_size,
+                padding=self.dw_padding,
+                bias=bias), BatchNorm2d(inner_channels // 4),
+            nn.ReLU(inplace=True),
+            DwPwConvTranspose(inner_channels // 4, inner_channels // 4, 2, 2),
+            BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),
+            DwPwConvTranspose(inner_channels // 4, 1, 2, 2), nn.Sigmoid())
+
+        self.adaptive = adaptive
+        if adaptive:
+            self.thresh = self._init_thresh(
+                inner_channels, serial=serial, smooth=smooth, bias=bias)
+
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1.)
+                m.bias.data.fill_(1e-4)
+
+    def _init_thresh(self,
+                     inner_channels,
+                     serial=False,
+                     smooth=False,
+                     bias=False):
+        in_channels = inner_channels
+        if serial:
+            in_channels += 1
+        self.thresh = nn.Sequential(
+            nn.Conv2d(
+                inner_channels,
+                inner_channels // 4,
+                self.dw_kernel_size,
+                padding=self.dw_padding,
+                bias=bias), BatchNorm2d(inner_channels // 4),
+            nn.ReLU(inplace=True),
+            self._init_upsample(
+                inner_channels // 4,
+                inner_channels // 4,
+                smooth=smooth,
+                bias=bias), BatchNorm2d(inner_channels // 4),
+            nn.ReLU(inplace=True),
+            self._init_upsample(
+                inner_channels // 4, 1, smooth=smooth, bias=bias),
+            nn.Sigmoid())
+
+        return self.thresh
+
+    def _init_upsample(self,
+                       in_channels,
+                       out_channels,
+                       smooth=False,
+                       bias=False):
+        if smooth:
+            inter_out_channels = out_channels
+            if out_channels == 1:
+                inter_out_channels = in_channels
+            module_list = [
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv2d(in_channels, inter_out_channels, 3, 1, 1, bias=bias)
+            ]
+            if out_channels == 1:
+                module_list.append(
+                    nn.Conv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        padding=1,
+                        bias=True))
+
+            return nn.Sequential(module_list)
+        else:
+            return nn.ConvTranspose2d(in_channels, out_channels, 2, 2)
+
+    def forward(self, features, gt=None, masks=None, training=False):
+        c2, c3, c4, c5 = features
+        p5 = self.up5(self.in5(c5))
+        p4 = self.up4(self.in4(c4))
+        p3 = self.up3(self.in3(c3))
+        p2 = self.in2(c2)
+
+        fuse = p5 + p4 + p3 + p2
+
+        # this is the pred module, not binarization module;
+        # We do not correct the name due to the trained model.
+        binary = self.binarize(fuse)
+        if self.training:
+            result = OrderedDict(binary=binary)
+        else:
+            return binary
+        if self.adaptive and self.training:
+            if self.serial:
+                fuse = torch.cat(
+                    (fuse, nn.functional.interpolate(binary, fuse.shape[2:])),
+                    1)
+            thresh = self.thresh(fuse)
+            thresh_binary = self.step_function(binary, thresh)
+            result.update(thresh=thresh, thresh_binary=thresh_binary)
+        return result
+
+    def step_function(self, x, y):
+        return torch.reciprocal(1 + torch.exp(-self.k * (x - y)))
+
+
 class SegDetector(nn.Module):
 
     def __init__(self,
@@ -471,6 +690,28 @@ class VLPTModel(nn.Module):
         return self.decoder(self.backbone(x))
 
 
+class DBNasModel(nn.Module):
+
+    def __init__(self, *args, **kwargs):
+        """
+        DB-NAS model
+        """
+        super(DBNasModel, self).__init__()
+        self.backbone = CompactDetBackbone(
+            width_stages=[32, 64, 96, 128], input_channel=32, **kwargs)
+        self.decoder = LightSegDetector(
+            in_channels=[32, 64, 96, 128],
+            adaptive=True,
+            k=50,
+            inner_channels=64,
+            dw_kernel_size=5,
+            dw_padding=2,
+            **kwargs)
+
+    def forward(self, x):
+        return self.decoder(self.backbone(x))
+
+
 class DBModel(nn.Module):
 
     def __init__(self, *args, **kwargs):

modelscope/models/cv/ocr_detection/modules/layers.py

@@ -0,0 +1,906 @@
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_same_padding(kernel_size):
+    if isinstance(kernel_size, tuple):
+        assert len(kernel_size) == 2, 'invalid kernel size: %s' % kernel_size
+        p1 = get_same_padding(kernel_size[0])
+        p2 = get_same_padding(kernel_size[1])
+        return p1, p2
+    assert isinstance(kernel_size,
+                      int), 'kernel size should be either `int` or `tuple`'
+    assert kernel_size % 2 > 0, 'kernel size should be odd number'
+    return kernel_size // 2
+
+
+def count_conv_flop(layer, x):
+    out_h = int(x.size(2) / layer.stride[0])
+    out_w = int(x.size(3) / layer.stride[1])
+    delta_ops = layer.in_channels * layer.out_channels * layer.kernel_size[
+        0] * layer.kernel_size[1] * out_h * out_w / layer.groups
+    return delta_ops
+
+
+def set_layer_from_config(layer_config):
+    if layer_config is None:
+        return None
+    name2layer = {
+        ZeroLayer.__name__: ZeroLayer,
+        MBInvertedConvLayer.__name__: MBInvertedConvLayer,
+        IdentityLayer.__name__: IdentityLayer,
+    }
+    layer_name = layer_config.pop('name')
+    layer = name2layer[layer_name]
+    return layer.build_from_config(layer_config)
+
+
+class MobileInvertedResidualBlock(nn.Module):
+
+    def __init__(self, mobile_inverted_conv, shortcut):
+        super(MobileInvertedResidualBlock, self).__init__()
+        self.mobile_inverted_conv = mobile_inverted_conv
+        self.shortcut = shortcut
+
+    def forward(self, x):
+        if self.mobile_inverted_conv.is_zero_layer():
+            res = x
+        elif self.shortcut is None or self.shortcut.is_zero_layer():
+            res = self.mobile_inverted_conv(x)
+        else:
+            conv_x = self.mobile_inverted_conv(x)
+            skip_x = self.shortcut(x)
+            res = skip_x + conv_x
+        return res
+
+    @property
+    def module_str(self):
+        return '(%s, %s)' % (self.mobile_inverted_conv.module_str,
+                             self.shortcut.module_str
+                             if self.shortcut is not None else None)
+
+    @property
+    def config(self):
+        return {
+            'name':
+            MobileInvertedResidualBlock.__name__,
+            'mobile_inverted_conv':
+            self.mobile_inverted_conv.config,
+            'shortcut':
+            self.shortcut.config if self.shortcut is not None else None,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        mobile_inverted_conv = set_layer_from_config(
+            config['mobile_inverted_conv'])
+        shortcut = set_layer_from_config(config['shortcut'])
+        return MobileInvertedResidualBlock(mobile_inverted_conv, shortcut)
+
+    def get_flops(self, x):
+        flops1, _ = self.mobile_inverted_conv.get_flops(x)
+        if self.shortcut:
+            flops2, _ = self.shortcut.get_flops(x)
+        else:
+            flops2 = 0
+        return flops1 + flops2, self.forward(x)
+
+
+class MBInvertedConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 stride=(1, 1),
+                 expand_ratio=6,
+                 mid_channels=None):
+        super(MBInvertedConvLayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+        if self.expand_ratio == 1:
+            self.inverted_bottleneck = None
+        else:
+            self.inverted_bottleneck = nn.Sequential(
+                OrderedDict([
+                    ('conv',
+                     nn.Conv2d(
+                         self.in_channels, feature_dim, 1, 1, 0, bias=False)),
+                    ('bn', nn.BatchNorm2d(feature_dim)),
+                    ('act', nn.PReLU()),
+                ]))
+        pad = get_same_padding(self.kernel_size)
+        self.depth_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(
+                     feature_dim,
+                     feature_dim,
+                     kernel_size,
+                     stride,
+                     pad,
+                     groups=feature_dim,
+                     bias=False)),
+                ('bn', nn.BatchNorm2d(feature_dim)),
+                ('act', nn.PReLU()),
+            ]))
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        if self.inverted_bottleneck:
+            x = self.inverted_bottleneck(x)
+        x = self.depth_conv(x)
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%dx%d_MBConv%d' % (self.kernel_size, self.kernel_size,
+                                   self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'kernel_size': self.kernel_size,
+            'stride': self.stride,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        if self.inverted_bottleneck:
+            total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+            x = self.inverted_bottleneck(x)
+        total_flops += count_conv_flop(self.depth_conv.conv, x)
+        x = self.depth_conv(x)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class IdentityLayer(nn.Module):
+
+    def __init__(self):
+        super(IdentityLayer, self).__init__()
+
+    def forward(self, x):
+        return x
+
+    @property
+    def module_str(self):
+        return 'Identity'
+
+    @property
+    def config(self):
+        return {
+            'name': IdentityLayer.__name__,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return IdentityLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        return 0, self.forward(x)
+
+
+class ZeroLayer(nn.Module):
+
+    def __init__(self, stride):
+        super(ZeroLayer, self).__init__()
+        self.stride = stride
+
+    def forward(self, x):
+        n, c, h, w = x.shape
+        h //= self.stride[0]
+        w //= self.stride[1]
+        device = x.device
+        padding = torch.zeros(n, c, h, w, device=device, requires_grad=False)
+        return padding
+
+    @property
+    def module_str(self):
+        return 'Zero'
+
+    @property
+    def config(self):
+        return {'name': ZeroLayer.__name__, 'stride': self.stride}
+
+    @staticmethod
+    def build_from_config(config):
+        return ZeroLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return True
+
+    def get_flops(self, x):
+        return 0, self.forward(x)
+
+
+def split_layer(total_channels, num_groups):
+    split = [
+        int(np.ceil(total_channels / num_groups)) for _ in range(num_groups)
+    ]
+    split[num_groups - 1] += total_channels - sum(split)
+    return split
+
+
+class MBInvertedMixConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 mix_conv_size=[1, 3, 5],
+                 stride=(1, 1),
+                 expand_ratio=6,
+                 mid_channels=None):
+        super(MBInvertedMixConvLayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.mix_conv_size = mix_conv_size
+        self.stride = stride
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+        self.inverted_bottleneck = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(self.in_channels, feature_dim, 1, 1, 0,
+                           bias=False)),
+                ('bn', nn.BatchNorm2d(feature_dim)),
+                ('act', nn.PReLU()),
+            ]))
+
+        self.mix_conv_size = mix_conv_size
+        self.n_chunks = len(mix_conv_size)
+        self.split_in_channels = split_layer(feature_dim, self.n_chunks)
+
+        self.mix_conv = nn.ModuleList()
+        for idx in range(self.n_chunks):
+            kernel_size = self.mix_conv_size[idx]
+            pad = get_same_padding(kernel_size)
+            split_in_channels_ = self.split_in_channels[idx]
+            self.mix_conv.append(
+                nn.Sequential(
+                    OrderedDict([
+                        ('conv',
+                         nn.Conv2d(
+                             split_in_channels_,
+                             split_in_channels_,
+                             kernel_size,
+                             stride,
+                             pad,
+                             groups=split_in_channels_,
+                             bias=False)),
+                        ('bn', nn.BatchNorm2d(split_in_channels_)),
+                        ('act', nn.PReLU()),
+                    ])))
+
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        x = self.inverted_bottleneck(x)
+        split = torch.split(x, self.split_in_channels, dim=1)
+        x = torch.cat([layer(s) for layer, s in zip(self.mix_conv, split)],
+                      dim=1)
+
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%s_MixConv%d' % (str(self.mix_conv_size), self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedMixConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'mix_conv_size': self.mix_conv_size,
+            'stride': self.stride,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedMixConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+        x = self.inverted_bottleneck(x)
+        split = torch.split(x, self.split_in_channels, dim=1)
+        out = []
+        for layer, s in zip(self.mix_conv, split):
+            out.append(layer(s))
+            total_flops += count_conv_flop(layer.conv, s)
+        x = torch.cat(out, dim=1)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class LinearMixConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 mix_conv_size=[1, 3, 5],
+                 stride=(1, 1),
+                 mid_channels=None):
+        super(LinearMixConvLayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.mix_conv_size = mix_conv_size
+        self.stride = stride
+        self.mid_channels = mid_channels
+
+        self.mix_conv_size = mix_conv_size
+        self.n_chunks = len(mix_conv_size)
+        self.mix_conv = nn.ModuleList()
+
+        for idx in range(self.n_chunks):
+            kernel_size = self.mix_conv_size[idx]
+            pad = get_same_padding(kernel_size)
+            self.mix_conv.append(
+                nn.Sequential(
+                    OrderedDict([
+                        ('conv',
+                         nn.Conv2d(
+                             in_channels,
+                             in_channels,
+                             kernel_size,
+                             stride,
+                             pad,
+                             groups=in_channels,
+                             bias=False)),
+                        ('bn', nn.BatchNorm2d(in_channels)),
+                        ('act', nn.ReLU6(inplace=True)),
+                    ])))
+
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(
+                     in_channels * self.n_chunks,
+                     out_channels,
+                     1,
+                     1,
+                     0,
+                     bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        x = torch.cat([layer(x) for layer in self.mix_conv], dim=1)
+
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%s_LinearMixConv' % (str(self.mix_conv_size))
+
+    @property
+    def config(self):
+        return {
+            'name': LinearMixConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'mix_conv_size': self.mix_conv_size,
+            'stride': self.stride,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return LinearMixConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        out = []
+        for layer in self.mix_conv:
+            out.append(layer(x))
+            total_flops += count_conv_flop(layer.conv, x)
+        x = torch.cat(out, dim=1)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class SELayer(nn.Module):
+    '''
+    '''
+
+    def __init__(self, input_channels: int, squeeze_factor: int = 4):
+        super(SELayer, self).__init__()
+        self.input_channels = input_channels
+        self.squeeze_factor = squeeze_factor
+        self.squeeze_channels = input_channels // squeeze_factor
+        self.fc1 = nn.Conv2d(self.input_channels, self.squeeze_channels, 1)
+        self.relu = nn.ReLU(inplace=True)
+        self.fc2 = nn.Conv2d(self.squeeze_channels, self.input_channels, 1)
+
+    def _scale(self, input):
+        scale = F.adaptive_avg_pool2d(input, 1)
+        scale = self.fc1(scale)
+        scale = self.relu(scale)
+        scale = self.fc2(scale)
+        return torch.sigmoid(scale)
+
+    def forward(self, input):
+        scale = self._scale(input)
+        return scale * input
+
+    @property
+    def module_str(self):
+        return 'SE_%d' % (self.squeeze_factor)
+
+    @property
+    def config(self):
+        return {
+            'name': SELayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'squeeze_factor': self.squeeze_factor,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return SELayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''
+        count se flops, only compute the fc layers' calculation
+        '''
+        total_flops = 0
+        total_flops += self.input_channels * self.squeeze_channels * 2
+        b, c, h, w = x.shape
+        total_flops += c * h * w
+        return total_flops, x
+
+
+class MHSA(nn.Module):
+
+    def __init__(self, n_dims, width=14, height=14, heads=4):
+        super(MHSA, self).__init__()
+        self.heads = heads
+
+        self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
+        self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
+        self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)
+
+        self.width = width
+        self.height = height
+
+        self.rel_h = nn.Parameter(
+            torch.randn([1, heads, n_dims // heads, 1, height]),
+            requires_grad=True)
+        self.rel_w = nn.Parameter(
+            torch.randn([1, heads, n_dims // heads, width, 1]),
+            requires_grad=True)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+
+        n_batch, C, width, height = x.size()
+        q = self.query(x).view(n_batch, self.heads, C // self.heads, -1)
+        k = self.key(x).view(n_batch, self.heads, C // self.heads, -1)
+        v = self.value(x).view(n_batch, self.heads, C // self.heads, -1)
+
+        content_content = torch.matmul(q.permute(0, 1, 3, 2), k)
+
+        content_position = (self.rel_h + self.rel_w).view(
+            1, self.heads, C // self.heads, -1).permute(0, 1, 3, 2)
+        content_position = torch.matmul(content_position, q)
+        energy = content_content + content_position
+        attention = self.softmax(energy)
+
+        out = torch.matmul(v, attention.permute(0, 1, 3, 2))
+        out = out.view(n_batch, C, width, height)
+
+        return out
+
+    def get_flops(self, x):
+        '''
+        count se flops, only compute the fc layers' calculation
+        '''
+        n_batch, C, width, height = x.size()
+
+        total_flops = 0
+        total_flops += count_conv_flop(self.query, x) * 3
+
+        # content_content
+        total_flops += (width * height) * C * (width * height)
+
+        # content_position
+        total_flops += (width * height) * C
+        total_flops += (width * height) * C * (width * height)
+
+        # attention
+        total_flops += (width * height) * C * (width * height)
+
+        return total_flops, x
+
+
+class MBInvertedMHSALayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 expand_ratio=6,
+                 width=1,
+                 height=175,
+                 mid_channels=None):
+        super(MBInvertedMHSALayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+        if self.expand_ratio == 1:
+            self.inverted_bottleneck = None
+        else:
+            self.inverted_bottleneck = nn.Sequential(
+                OrderedDict([
+                    ('conv',
+                     nn.Conv2d(
+                         self.in_channels, feature_dim, 1, 1, 0, bias=False)),
+                    ('bn', nn.BatchNorm2d(feature_dim)),
+                    #                 ('act', nn.PReLU()),
+                    ('act', nn.ReLU6(inplace=True)),
+                ]))
+        self.mhsa = MHSA(feature_dim, width, height)
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        if self.inverted_bottleneck:
+            x = self.inverted_bottleneck(x)
+        x = self.mhsa(x)
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return 'MSHA%d' % (self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedMHSALayer.__name__,
+            'in_channels': self.in_channels,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedMHSALayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        if self.inverted_bottleneck:
+            total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+            x = self.inverted_bottleneck(x)
+        total_flops += self.mhsa.get_flops(x)[0]
+        x = self.mhsa(x)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class MBInvertedRepConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 rep_conv_size=[1, 3, 5],
+                 stride=(1, 1),
+                 expand_ratio=6,
+                 mid_channels=None):
+        super(MBInvertedRepConvLayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.rep_conv_size = rep_conv_size
+        self.stride = stride
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+
+        self.inverted_bottleneck = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(self.in_channels, feature_dim, 1, 1, 0,
+                           bias=False)),
+                ('bn', nn.BatchNorm2d(feature_dim)),
+                ('act', nn.PReLU()),
+            ]))
+
+        self.rep_conv_size = rep_conv_size
+        self.n_chunks = len(rep_conv_size)
+
+        self.rep_conv = nn.ModuleList()
+        for idx in range(self.n_chunks):
+            kernel_size = self.rep_conv_size[idx]
+            pad = get_same_padding(kernel_size)
+            self.rep_conv.append(
+                nn.Sequential(
+                    OrderedDict([
+                        ('conv',
+                         nn.Conv2d(
+                             feature_dim,
+                             feature_dim,
+                             kernel_size,
+                             stride,
+                             pad,
+                             groups=feature_dim,
+                             bias=False)),
+                        ('bn', nn.BatchNorm2d(feature_dim)),
+                    ])))
+
+        self.rep_conv_deploy = None
+        self.act = nn.PReLU()
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+        self.deploy = False
+
+    def forward(self, x):
+        x = self.inverted_bottleneck(x)
+        if not self.deploy:
+            out = []
+            for layer in self.rep_conv:
+                out.append(layer(x))
+            x = out[0]
+            for out_ in out[1:]:
+                x += out_
+        else:
+            x = self.rep_conv_deploy(x)
+        x = self.act(x)
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%s_RepConv%d' % (str(self.rep_conv_size), self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedMixConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'rep_conv_size': self.rep_conv_size,
+            'stride': self.stride,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedMixConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def switch_to_deploy(self):
+        self.deploy = True
+        feature_dim = self.rep_conv[0].conv.in_channels
+        stride = self.rep_conv[0].conv.stride
+        kernel_size = max(self.rep_conv_size)
+        pad = get_same_padding(kernel_size)
+        self.rep_conv_deploy = nn.Conv2d(
+            feature_dim,
+            feature_dim,
+            kernel_size,
+            stride,
+            pad,
+            groups=feature_dim,
+            bias=True)
+
+        kernel, bias = self.get_equivalent_kernel_bias()
+
+        self.rep_conv_deploy.weight.data = kernel
+        self.rep_conv_deploy.bias.data = bias
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__('rep_conv')
+
+    def get_equivalent_kernel_bias(self):
+        max_kernel_size = max(self.rep_conv_size)
+
+        if max_kernel_size == 5:
+            if 1 in self.rep_conv_size:
+                kernel1x1, bias1x1 = self._fuse_bn_tensor(
+                    self.rep_conv[self.rep_conv_size.index(1)])
+            else:
+                kernel1x1 = None
+                bias1x1 = 0
+
+            if 3 in self.rep_conv_size:
+                kernel3x3, bias3x3 = self._fuse_bn_tensor(
+                    self.rep_conv[self.rep_conv_size.index(3)])
+            else:
+                kernel3x3 = None
+                bias3x3 = 0
+
+            if 5 in self.rep_conv_size:
+                kernel5x5, bias5x5 = self._fuse_bn_tensor(
+                    self.rep_conv[self.rep_conv_size.index(5)])
+
+            else:
+                kernel5x5 = 0
+                bias5x5 = 0
+
+            return kernel5x5 + self._pad_1x1_to_5x5_tensor(
+                kernel1x1) + self._pad_3x3_to_5x5_tensor(
+                    kernel3x3), bias5x5 + bias3x3 + bias1x1
+        else:
+            if 1 in self.rep_conv_size:
+                kernel1x1, bias1x1 = self._fuse_bn_tensor(
+                    self.rep_conv[self.rep_conv_size.index(1)])
+            else:
+                kernel1x1 = None
+                bias1x1 = 0
+
+            if 3 in self.rep_conv_size:
+                kernel3x3, bias3x3 = self._fuse_bn_tensor(
+                    self.rep_conv[self.rep_conv_size.index(3)])
+            else:
+                kernel3x3 = None
+                bias3x3 = 0
+
+            return kernel3x3 + self._pad_1x1_to_3x3_tensor(
+                kernel1x1), bias3x3 + bias1x1
+
+    def _pad_1x1_to_3x3_tensor(self, kernel1x1):
+        if kernel1x1 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])
+
+    def _pad_1x1_to_5x5_tensor(self, kernel1x1):
+        if kernel1x1 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel1x1, [2, 2, 2, 2])
+
+    def _pad_3x3_to_5x5_tensor(self, kernel3x3):
+        if kernel3x3 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel3x3, [1, 1, 1, 1])
+
+    def _fuse_bn_tensor(self, branch):
+        if branch is None:
+            return 0, 0
+        if isinstance(branch, nn.Sequential):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, 'id_tensor'):
+                input_dim = self.in_channels // self.groups
+                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3),
+                                        dtype=np.float32)
+                for i in range(self.in_channels):
+                    kernel_value[i, i % input_dim, 1, 1] = 1
+                self.id_tensor = torch.from_numpy(kernel_value).to(
+                    branch.weight.device)
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+
+        return kernel * t, beta - running_mean * gamma / std
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+        x = self.inverted_bottleneck(x)
+
+        total_flops += count_conv_flop(self.rep_conv[-1].conv, x)
+        out = []
+        for layer in self.rep_conv:
+            out.append(layer(x))
+        x = out[0]
+        for out_ in out[1:]:
+            x += out_
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x

modelscope/models/cv/ocr_detection/modules/mix_ops.py

@@ -0,0 +1,689 @@
+# ------------------------------------------------------------------------------
+# Part of implementation is adopted from ProxylessNAS,
+# made publicly available under the Apache License 2.0 at https://github.com/mit-han-lab/proxylessnas.
+# ------------------------------------------------------------------------------
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+
+from .layers import (IdentityLayer, LinearMixConvLayer, MBInvertedConvLayer,
+                     MBInvertedMHSALayer, MBInvertedMixConvLayer,
+                     MBInvertedRepConvLayer, SELayer, ZeroLayer)
+
+
+def detach_variable(inputs):
+    if isinstance(inputs, tuple):
+        return tuple([detach_variable(x) for x in inputs])
+    else:
+        x = inputs.detach()
+        x.requires_grad = inputs.requires_grad
+        return x
+
+
+def delta_ij(i, j):
+    if i == j:
+        return 1
+    else:
+        return 0
+
+
+def conv_func_by_name(name):
+    name2ops = {
+        'Identity':
+        lambda in_C, out_C, S: IdentityLayer(in_C, out_C, ops_order=ops_order),
+        'Zero':
+        lambda in_C, out_C, S: ZeroLayer(stride=S),
+    }
+    name2ops.update({
+        '3x3_MBConv1':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 1),
+        '3x3_MBConv2':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 2),
+        '3x3_MBConv3':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 3),
+        '3x3_MBConv4':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 4),
+        '3x3_MBConv5':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 5),
+        '3x3_MBConv6':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 6),
+        #######################################################################################
+        '5x5_MBConv1':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 1),
+        '5x5_MBConv2':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 2),
+        '5x5_MBConv3':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 3),
+        '5x5_MBConv4':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 4),
+        '5x5_MBConv5':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 5),
+        '5x5_MBConv6':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 6),
+        #######################################################################################
+        '7x7_MBConv1':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 1),
+        '7x7_MBConv2':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 2),
+        '7x7_MBConv3':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 3),
+        '7x7_MBConv4':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 4),
+        '7x7_MBConv5':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 5),
+        '7x7_MBConv6':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 6),
+        #######################################################################################
+        '13_MixConv1':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 1
+                                                      ),
+        '13_MixConv2':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 2
+                                                      ),
+        '13_MixConv3':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 3
+                                                      ),
+        '13_MixConv4':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 4
+                                                      ),
+        '13_MixConv5':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 5
+                                                      ),
+        '13_MixConv6':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 6
+                                                      ),
+        #######################################################################################
+        '35_MixConv1':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 1
+                                                      ),
+        '35_MixConv2':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 2
+                                                      ),
+        '35_MixConv3':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 3
+                                                      ),
+        '35_MixConv4':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 4
+                                                      ),
+        '35_MixConv5':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 5
+                                                      ),
+        '35_MixConv6':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 6
+                                                      ),
+        #######################################################################################
+        '135_MixConv1':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 1),
+        '135_MixConv2':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 2),
+        '135_MixConv3':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 3),
+        '135_MixConv4':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 4),
+        '135_MixConv5':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 5),
+        '135_MixConv6':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 6),
+        #######################################################################################
+        '13_LinMixConv':
+        lambda in_C, out_C, S: LinearMixConvLayer(in_C, out_C, [1, 3], S),
+        '35_LinMixConv':
+        lambda in_C, out_C, S: LinearMixConvLayer(in_C, out_C, [3, 5], S),
+        '135_LinMixConv':
+        lambda in_C, out_C, S: LinearMixConvLayer(in_C, out_C, [1, 3, 5], S),
+        #######################################################################################
+        'SE_2':
+        lambda in_C, out_C, S: SELayer(in_C, 2),
+        'SE_4':
+        lambda in_C, out_C, S: SELayer(in_C, 4),
+        'SE_8':
+        lambda in_C, out_C, S: SELayer(in_C, 8),
+        #######################################################################################
+        'MHSA1':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 1, width, height),
+        'MHSA2':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 2, width, height),
+        'MHSA3':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 3, width, height),
+        'MHSA4':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 4, width, height),
+        'MHSA5':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 5, width, height),
+        'MHSA6':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 6, width, height),
+        #######################################################################################
+        '13_RepConv1':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 1
+                                                      ),
+        '13_RepConv2':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 2
+                                                      ),
+        '13_RepConv3':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 3
+                                                      ),
+        '13_RepConv4':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 4
+                                                      ),
+        '13_RepConv5':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 5
+                                                      ),
+        '13_RepConv6':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 6
+                                                      ),
+        #######################################################################################
+        '35_RepConv1':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 1
+                                                      ),
+        '35_RepConv2':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 2
+                                                      ),
+        '35_RepConv3':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 3
+                                                      ),
+        '35_RepConv4':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 4
+                                                      ),
+        '35_RepConv5':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 5
+                                                      ),
+        '35_RepConv6':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 6
+                                                      ),
+        #######################################################################################
+        '135_RepConv1':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 1),
+        '135_RepConv2':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 2),
+        '135_RepConv3':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 3),
+        '135_RepConv4':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 4),
+        '135_RepConv5':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 5),
+        '135_RepConv6':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 6),
+    })
+    return name2ops[name]
+
+
+def build_candidate_ops(candidate_ops,
+                        in_channels,
+                        out_channels,
+                        stride,
+                        ops_order,
+                        spatio_size=None):
+    if candidate_ops is None:
+        raise ValueError('please specify a candidate set')
+
+    name2ops = {
+        'Identity':
+        lambda in_C, out_C, S: IdentityLayer(in_C, out_C, ops_order=ops_order),
+        'Zero':
+        lambda in_C, out_C, S: ZeroLayer(stride=S),
+    }
+    # add MBConv layers
+    name2ops.update({
+        '3x3_MBConv1':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 1),
+        '3x3_MBConv2':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 2),
+        '3x3_MBConv3':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 3),
+        '3x3_MBConv4':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 4),
+        '3x3_MBConv5':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 5),
+        '3x3_MBConv6':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 3, S, 6),
+        #######################################################################################
+        '5x5_MBConv1':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 1),
+        '5x5_MBConv2':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 2),
+        '5x5_MBConv3':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 3),
+        '5x5_MBConv4':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 4),
+        '5x5_MBConv5':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 5),
+        '5x5_MBConv6':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 5, S, 6),
+        #######################################################################################
+        '7x7_MBConv1':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 1),
+        '7x7_MBConv2':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 2),
+        '7x7_MBConv3':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 3),
+        '7x7_MBConv4':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 4),
+        '7x7_MBConv5':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 5),
+        '7x7_MBConv6':
+        lambda in_C, out_C, S: MBInvertedConvLayer(in_C, out_C, 7, S, 6),
+        #######################################################################################
+        '13_MixConv1':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 1
+                                                      ),
+        '13_MixConv2':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 2
+                                                      ),
+        '13_MixConv3':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 3
+                                                      ),
+        '13_MixConv4':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 4
+                                                      ),
+        '13_MixConv5':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 5
+                                                      ),
+        '13_MixConv6':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3], S, 6
+                                                      ),
+        #######################################################################################
+        '35_MixConv1':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 1
+                                                      ),
+        '35_MixConv2':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 2
+                                                      ),
+        '35_MixConv3':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 3
+                                                      ),
+        '35_MixConv4':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 4
+                                                      ),
+        '35_MixConv5':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 5
+                                                      ),
+        '35_MixConv6':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [3, 5], S, 6
+                                                      ),
+        #######################################################################################
+        '135_MixConv1':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 1),
+        '135_MixConv2':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 2),
+        '135_MixConv3':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 3),
+        '135_MixConv4':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 4),
+        '135_MixConv5':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 5),
+        '135_MixConv6':
+        lambda in_C, out_C, S: MBInvertedMixConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 6),
+        #######################################################################################
+        '13_LinMixConv':
+        lambda in_C, out_C, S: LinearMixConvLayer(in_C, out_C, [1, 3], S),
+        '35_LinMixConv':
+        lambda in_C, out_C, S: LinearMixConvLayer(in_C, out_C, [3, 5], S),
+        '135_LinMixConv':
+        lambda in_C, out_C, S: LinearMixConvLayer(in_C, out_C, [1, 3, 5], S),
+        #######################################################################################
+        'SE_2':
+        lambda in_C, out_C, S: SELayer(in_C, 2),
+        'SE_4':
+        lambda in_C, out_C, S: SELayer(in_C, 4),
+        'SE_8':
+        lambda in_C, out_C, S: SELayer(in_C, 8),
+        #######################################################################################
+        'MHSA1':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 1, width, height),
+        'MHSA2':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 2, width, height),
+        'MHSA3':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 3, width, height),
+        'MHSA4':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 4, width, height),
+        'MHSA5':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 5, width, height),
+        'MHSA6':
+        lambda in_C, out_C, width, height: MBInvertedMHSALayer(
+            in_C, out_C, 6, width, height),
+        #######################################################################################
+        '13_RepConv1':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 1
+                                                      ),
+        '13_RepConv2':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 2
+                                                      ),
+        '13_RepConv3':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 3
+                                                      ),
+        '13_RepConv4':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 4
+                                                      ),
+        '13_RepConv5':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 5
+                                                      ),
+        '13_RepConv6':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3], S, 6
+                                                      ),
+        #######################################################################################
+        '35_RepConv1':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 1
+                                                      ),
+        '35_RepConv2':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 2
+                                                      ),
+        '35_RepConv3':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 3
+                                                      ),
+        '35_RepConv4':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 4
+                                                      ),
+        '35_RepConv5':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 5
+                                                      ),
+        '35_RepConv6':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [3, 5], S, 6
+                                                      ),
+        #######################################################################################
+        '135_RepConv1':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 1),
+        '135_RepConv2':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 2),
+        '135_RepConv3':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 3),
+        '135_RepConv4':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 4),
+        '135_RepConv5':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 5),
+        '135_RepConv6':
+        lambda in_C, out_C, S: MBInvertedRepConvLayer(in_C, out_C, [1, 3, 5],
+                                                      S, 6),
+    })
+
+    out = []
+    for name in candidate_ops:
+
+        if 'MHSA' in name:
+            out.append(name2ops[name](in_channels, out_channels,
+                                      spatio_size[0], spatio_size[1]))
+        else:
+            out.append(name2ops[name](in_channels, out_channels, stride))
+    return out
+
+
+class MixedEdge(nn.Module):
+    MODE = None  # full, two, None, full_v2
+
+    def __init__(self, candidate_ops):
+        super(MixedEdge, self).__init__()
+
+        self.candidate_ops = nn.ModuleList(candidate_ops)
+        self.AP_path_alpha = Parameter(torch.Tensor(
+            self.n_choices))  # architecture parameters
+        self.AP_path_wb = Parameter(torch.Tensor(
+            self.n_choices))  # binary gates
+
+        self.active_index = [0]
+        self.inactive_index = None
+
+        self.log_prob = None
+        self.current_prob_over_ops = None
+
+    @property
+    def n_choices(self):
+        return len(self.candidate_ops)
+
+    @property
+    def probs_over_ops(self):
+        probs = F.softmax(self.AP_path_alpha, dim=0)  # softmax to probability
+        return probs
+
+    @property
+    def chosen_index(self):
+        probs = self.probs_over_ops.data.cpu().numpy()
+        index = int(np.argmax(probs))
+        return index, probs[index]
+
+    @property
+    def chosen_op(self):
+        index, _ = self.chosen_index
+        return self.candidate_ops[index]
+
+    @property
+    def random_op(self):
+        index = np.random.choice([_i for _i in range(self.n_choices)], 1)[0]
+        return self.candidate_ops[index]
+
+    def entropy(self, eps=1e-8):
+        probs = self.probs_over_ops
+        log_probs = torch.log(probs + eps)
+        entropy = -torch.sum(torch.mul(probs, log_probs))
+        return entropy
+
+    def is_zero_layer(self):
+        return self.active_op.is_zero_layer()
+
+    @property
+    def active_op(self):
+        """ assume only one path is active """
+        return self.candidate_ops[self.active_index[0]]
+
+    def set_chosen_op_active(self):
+        chosen_idx, _ = self.chosen_index
+        self.active_index = [chosen_idx]
+        self.inactive_index = [_i for _i in range(0, chosen_idx)] + \
+                              [_i for _i in range(chosen_idx + 1, self.n_choices)]
+
+    def forward(self, x):
+        if MixedEdge.MODE == 'full' or MixedEdge.MODE == 'two':
+            output = 0
+            for _i in self.active_index:
+                oi = self.candidate_ops[_i](x)
+                output = output + self.AP_path_wb[_i] * oi
+            for _i in self.inactive_index:
+                oi = self.candidate_ops[_i](x)
+                output = output + self.AP_path_wb[_i] * oi.detach()
+        elif MixedEdge.MODE == 'full_v2':
+
+            def run_function(candidate_ops, active_id):
+
+                def forward(_x):
+                    return candidate_ops[active_id](_x)
+
+                return forward
+
+            def backward_function(candidate_ops, active_id, binary_gates):
+
+                def backward(_x, _output, grad_output):
+                    binary_grads = torch.zeros_like(binary_gates.data)
+                    with torch.no_grad():
+                        for k in range(len(candidate_ops)):
+                            if k != active_id:
+                                out_k = candidate_ops[k](_x.data)
+                            else:
+                                out_k = _output.data
+                            grad_k = torch.sum(out_k * grad_output)
+                            binary_grads[k] = grad_k
+                    return binary_grads
+
+                return backward
+
+            output = ArchGradientFunction.apply(
+                x, self.AP_path_wb,
+                run_function(self.candidate_ops, self.active_index[0]),
+                backward_function(self.candidate_ops, self.active_index[0],
+                                  self.AP_path_wb))
+        else:
+            output = self.active_op(x)
+        return output
+
+    @property
+    def module_str(self):
+        chosen_index, probs = self.chosen_index
+        return 'Mix(%s, %.3f)' % (self.candidate_ops[chosen_index].module_str,
+                                  probs)
+
+    @property
+    def config(self):
+        raise ValueError('not needed')
+
+    @staticmethod
+    def build_from_config(config):
+        raise ValueError('not needed')
+
+    def get_flops(self, x):
+        """ Only active paths taken into consideration when calculating FLOPs """
+        flops = 0
+        for i in self.active_index:
+            delta_flop, _ = self.candidate_ops[i].get_flops(x)
+            flops += delta_flop
+        return flops, self.forward(x)
+
+    def binarize(self):
+        """ prepare: active_index, inactive_index, AP_path_wb, log_prob (optional), current_prob_over_ops (optional) """
+        self.log_prob = None
+        # reset binary gates
+        self.AP_path_wb.data.zero_()
+        # binarize according to probs
+        probs = self.probs_over_ops
+        if MixedEdge.MODE == 'two':
+            # sample two ops according to `probs`
+            sample_op = torch.multinomial(probs.data, 2, replacement=False)
+            probs_slice = F.softmax(
+                torch.stack([self.AP_path_alpha[idx] for idx in sample_op]),
+                dim=0)
+            self.current_prob_over_ops = torch.zeros_like(probs)
+            for i, idx in enumerate(sample_op):
+                self.current_prob_over_ops[idx] = probs_slice[i]
+            # chose one to be active and the other to be inactive according to probs_slice
+            c = torch.multinomial(probs_slice.data, 1)[0]  # 0 or 1
+            active_op = sample_op[c].item()
+            inactive_op = sample_op[1 - c].item()
+            self.active_index = [active_op]
+            self.inactive_index = [inactive_op]
+            # set binary gate
+            self.AP_path_wb.data[active_op] = 1.0
+        else:
+            sample = torch.multinomial(probs.data, 1)[0].item()
+            self.active_index = [sample]
+            self.inactive_index = [_i for _i in range(0, sample)] + \
+                                  [_i for _i in range(sample + 1, self.n_choices)]
+            self.log_prob = torch.log(probs[sample])
+            self.current_prob_over_ops = probs
+            # set binary gate
+            self.AP_path_wb.data[sample] = 1.0
+        # avoid over-regularization
+        for _i in range(self.n_choices):
+            for name, param in self.candidate_ops[_i].named_parameters():
+                param.grad = None
+
+    def set_arch_param_grad(self):
+        binary_grads = self.AP_path_wb.grad.data
+        if self.active_op.is_zero_layer():
+            self.AP_path_alpha.grad = None
+            return
+        if self.AP_path_alpha.grad is None:
+            self.AP_path_alpha.grad = torch.zeros_like(self.AP_path_alpha.data)
+        if MixedEdge.MODE == 'two':
+            involved_idx = self.active_index + self.inactive_index
+            probs_slice = F.softmax(
+                torch.stack([self.AP_path_alpha[idx] for idx in involved_idx]),
+                dim=0).data
+            for i in range(2):
+                for j in range(2):
+                    origin_i = involved_idx[i]
+                    origin_j = involved_idx[j]
+                    self.AP_path_alpha.grad.data[origin_i] += \
+                        binary_grads[origin_j] * probs_slice[j] * (delta_ij(i, j) - probs_slice[i])
+            for _i, idx in enumerate(self.active_index):
+                self.active_index[_i] = (idx,
+                                         self.AP_path_alpha.data[idx].item())
+            for _i, idx in enumerate(self.inactive_index):
+                self.inactive_index[_i] = (idx,
+                                           self.AP_path_alpha.data[idx].item())
+        else:
+            probs = self.probs_over_ops.data
+            for i in range(self.n_choices):
+                for j in range(self.n_choices):
+                    self.AP_path_alpha.grad.data[
+                        i] += binary_grads[j] * probs[j] * (
+                            delta_ij(i, j) - probs[i])
+        return
+
+    def rescale_updated_arch_param(self):
+        if not isinstance(self.active_index[0], tuple):
+            assert self.active_op.is_zero_layer()
+            return
+        involved_idx = [
+            idx for idx, _ in (self.active_index + self.inactive_index)
+        ]
+        old_alphas = [
+            alpha for _, alpha in (self.active_index + self.inactive_index)
+        ]
+        new_alphas = [self.AP_path_alpha.data[idx] for idx in involved_idx]
+
+        offset = math.log(
+            sum([math.exp(alpha) for alpha in new_alphas])
+            / sum([math.exp(alpha) for alpha in old_alphas]))
+
+        for idx in involved_idx:
+            self.AP_path_alpha.data[idx] -= offset
+
+
+class ArchGradientFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, binary_gates, run_func, backward_func):
+        ctx.run_func = run_func
+        ctx.backward_func = backward_func
+
+        detached_x = detach_variable(x)
+        with torch.enable_grad():
+            output = run_func(detached_x)
+        ctx.save_for_backward(detached_x, output)
+        return output.data
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        detached_x, output = ctx.saved_tensors
+
+        grad_x = torch.autograd.grad(
+            output, detached_x, grad_output, only_inputs=True)
+        # compute gradients w.r.t. binary_gates
+        binary_grads = ctx.backward_func(detached_x.data, output.data,
+                                         grad_output.data)
+
+        return grad_x[0], binary_grads, None, None

modelscope/models/cv/ocr_detection/modules/proxyless.py

@@ -0,0 +1,178 @@
+import re
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .layers import (IdentityLayer, MBInvertedConvLayer,
+                     MobileInvertedResidualBlock, ZeroLayer)
+from .mix_ops import MixedEdge, build_candidate_ops, conv_func_by_name
+
+
+class NasRecBackbone(nn.Module):
+
+    def __init__(self, first_conv, blocks):
+        super(NasRecBackbone, self).__init__()
+        self.first_conv = first_conv
+        self.blocks = nn.ModuleList(blocks)
+        self.output_idx = [5, 11, 17, 23]
+
+    def forward(self, x):
+        x = self.first_conv(x)
+        idx = 0
+        out = []
+        for block in self.blocks:
+            x = block(x)
+            if (idx + 1) % (int(len(self.blocks) / 4)) == 0:
+                out.append(x)
+            idx += 1
+        return out[0], out[1], out[2], out[3]
+
+    def get_bn_param(self):
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
+                return {
+                    'momentum': m.momentum,
+                    'eps': m.eps,
+                }
+        return None
+
+    @property
+    def config(self):
+        return {
+            'name': NasRecBackbone.__name__,
+            'bn': self.get_bn_param(),
+            'first_conv': 'conv_in3_out32_k3_s2_p1',
+            'blocks': [block.config for block in self.blocks]
+        }
+
+    def set_bn_param(self, momentum, eps):
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
+                m.momentum = momentum
+                m.eps = eps
+        return
+
+    @staticmethod
+    def build_from_config(config):
+        first_conv_config = config['first_conv']
+        match_obj = re.match(r'conv_in(\d+)_out(\d+)_k(\d+)_s(\d+)_p(\d+)',
+                             first_conv_config)
+        in_channel = int(match_obj.group(1))
+        out_channel = int(match_obj.group(2))
+        kernel_size = int(match_obj.group(3))
+        stride = int(match_obj.group(4))
+        padding = int(match_obj.group(5))
+        first_conv = nn.Sequential(
+            nn.Conv2d(
+                in_channel,
+                out_channel,
+                kernel_size,
+                stride,
+                padding,
+                bias=False),
+            nn.BatchNorm2d(out_channel),
+            nn.ReLU(inplace=True),
+        )
+        blocks = []
+        for block_config in config['blocks']:
+            blocks.append(
+                MobileInvertedResidualBlock.build_from_config(block_config))
+        net = NasRecBackbone(first_conv, blocks)
+        if 'bn' in config:
+            net.set_bn_param(**config['bn'])
+        else:
+            net.set_bn_param(momentum=0.1, eps=1e-3)
+        return net
+
+
+class CompactDetBackbone(NasRecBackbone):
+    '''
+    proxyless nas backbone, 5M.
+    '''
+
+    def __init__(self,
+                 width_stages,
+                 input_channel=None,
+                 bn_param=(0.1, 1e-3),
+                 **kwargs):
+
+        if input_channel is None:
+            input_channel = width_stages[0]
+        first_conv = nn.Sequential(
+            nn.Conv2d(
+                3,
+                input_channel,
+                kernel_size=(3, 3),
+                stride=2,
+                padding=1,
+                bias=False), nn.BatchNorm2d(input_channel),
+            nn.ReLU(inplace=True))
+
+        conv_candidates = [
+            '5x5_MBConv2', '5x5_MBConv4', '3x3_MBConv2', '3x3_MBConv4',
+            '13_MixConv2', '13_MixConv4', '35_MixConv2', '35_MixConv4',
+            '135_MixConv2', '135_MixConv4', '13_LinMixConv', '35_LinMixConv',
+            '135_LinMixConv', '13_RepConv2', '13_RepConv4', '35_RepConv2',
+            '35_RepConv4', '135_RepConv2', '135_RepConv4', 'Zero'
+        ]
+
+        se_candidates = ['SE_2', 'SE_4', 'SE_8', 'Zero']
+
+        conv_op_ids = [
+            15, 17, 17, 17, 17, 0, 16, 16, 18, 18, 16, 2, 16, 18, 16, 18, 18,
+            2, 1, 18, 18, 18, 16, 2
+        ]
+
+        n_cell_stages = [5, 5, 5, 5]
+
+        stride_stages = [(2, 2), (2, 2), (2, 2), (2, 2)]
+
+        if se_candidates:
+            assert len(conv_op_ids) == sum(n_cell_stages) + 4
+        else:
+            assert len(conv_op_ids) == sum(n_cell_stages)
+
+        blocks = []
+
+        for width, n_cell, s in zip(width_stages, n_cell_stages,
+                                    stride_stages):
+            for i in range(n_cell):
+                if i == 0:
+                    stride = s
+                else:
+                    stride = (1, 1)
+                block_i = len(blocks)
+                conv_op = conv_func_by_name(
+                    conv_candidates[conv_op_ids[block_i]])(input_channel,
+                                                           width, stride)
+
+                if stride == (1, 1) and input_channel == width:
+                    shortcut = IdentityLayer()
+                else:
+                    shortcut = None
+                inverted_residual_block = MobileInvertedResidualBlock(
+                    conv_op, shortcut)
+                blocks.append(inverted_residual_block)
+                input_channel = width
+
+            if se_candidates is not None:
+                block_i = len(blocks)
+                se_op = conv_func_by_name(se_candidates[conv_op_ids[block_i]])(
+                    input_channel, width, stride)
+                shortcut = IdentityLayer()
+                inverted_residual_block = MobileInvertedResidualBlock(
+                    se_op, shortcut)
+                blocks.append(inverted_residual_block)
+
+        self.out_channel = input_channel
+
+        super(CompactDetBackbone, self).__init__(first_conv, blocks)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)

modelscope/models/cv/ocr_recognition/model.py

@@ -10,8 +10,9 @@ from modelscope.models.builder import MODELS
 from modelscope.utils.config import Config
 from modelscope.utils.constant import ModelFile, Tasks
 from modelscope.utils.logger import get_logger
-from .modules.convnextvit import ConvNextViT
-from .modules.crnn import CRNN
+from .modules.ConvNextViT.main_model import ConvNextViT
+from .modules.CRNN.main_model import CRNN
+from .modules.LightweightEdge.main_model import LightweightEdge
 
 LOGGER = get_logger()
 
@@ -85,6 +86,8 @@ class OCRRecognition(TorchModel):
             self.recognizer = ConvNextViT()
         elif cfgs.model.recognizer == 'CRNN':
             self.recognizer = CRNN()
+        elif cfgs.model.recognizer == 'LightweightEdge':
+            self.recognizer = LightweightEdge()
         else:
             raise TypeError(
                 f'recognizer should be either ConvNextViT, CRNN, but got {cfgs.model.recognizer}'
@@ -94,7 +97,7 @@ class OCRRecognition(TorchModel):
             model_dict = self.recognizer.state_dict()
             # remove prefix for finetuned models
             check_point = {
-                k.replace('recognizer.', ''): v
+                k.replace('recognizer.', '').replace('module.', ''): v
                 for k, v in params_pretrained.items()
             }
             model_dict.update(check_point)
@@ -134,9 +137,9 @@ class OCRRecognition(TorchModel):
         labels = batch['labels']
         bs = inputs.shape[0]
         if self.do_chunking:
-            inputs = inputs.view(bs * 3, 1, self.target_height, 300)
+            inputs = inputs.view(bs * 3, 3, self.target_height, 300)
         else:
-            inputs = inputs.view(bs, 1, self.target_height, self.target_width)
+            inputs = inputs.view(bs, 3, self.target_height, self.target_width)
         output = self(inputs)
         probs = output['probs'].permute(1, 0, 2)
         _, label_length, label_flatten = self.encdec.encode(labels)

modelscope/models/cv/ocr_recognition/modules/CRNN/main_model.py

@@ -79,6 +79,11 @@ class CRNN(nn.Module):
         self.cls = nn.Linear(512, 7644, bias=False)
 
     def forward(self, input):
+        # RGB2GRAY
+        input = input[:, 0:
+                      1, :, :] * 0.2989 + input[:, 1:
+                                                2, :, :] * 0.5870 + input[:, 2:
+                                                                          3, :, :] * 0.1140
         feats = self.conv0(input)
         feats = self.p0(feats)
         feats = self.conv1(feats)

modelscope/models/cv/ocr_recognition/modules/ConvNextViT/main_model.py

@@ -14,7 +14,11 @@ class ConvNextViT(nn.Module):
         self.vitstr = vitstr_tiny(num_tokens=7644)
 
     def forward(self, input):
-        """ Transformation stage """
+        # RGB2GRAY
+        input = input[:, 0:
+                      1, :, :] * 0.2989 + input[:, 1:
+                                                2, :, :] * 0.5870 + input[:, 2:
+                                                                          3, :, :] * 0.1140
         features = self.cnn_model(input)
         output = self.vitstr(features)
         return output

modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py

@@ -0,0 +1,42 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+from collections import OrderedDict
+
+import torch.nn as nn
+
+from .nas_block import plnas_linear_mix_se
+
+
+class LightweightEdge(nn.Module):
+    """
+        基于混合rep block的nas模型
+        Args:
+            input (tensor): batch of input images
+    """
+
+    def __init__(self):
+        super(LightweightEdge, self).__init__()
+        self.FeatureExtraction = plnas_linear_mix_se(3, 123)
+        self.AdaptiveAvgPool = nn.AdaptiveAvgPool2d(
+            (None, 1))  # Transform final (imgH/16-1) -> 1
+        self.dropout = nn.Dropout(0.3)
+        self.Prediction = nn.Sequential(
+            OrderedDict([
+                ('fc1', nn.Linear(123, 120)),
+                ('bn', nn.BatchNorm1d(120)),
+                ('fc2', nn.Linear(120, 7642)),
+            ]))
+
+    def forward(self, input):
+        visual_feature = self.FeatureExtraction(input)
+        visual_feature = self.AdaptiveAvgPool(
+            visual_feature.permute(0, 3, 1, 2))  # [b, c, h, w] -> [b, w, c, h]
+        visual_feature = visual_feature.squeeze(3)
+        visual_feature = self.dropout(visual_feature)
+        prediction = self.Prediction.fc1(visual_feature.contiguous())
+        b, t, c = prediction.shape
+        prediction = self.Prediction.bn(prediction.view(b * t,
+                                                        c)).view(b, t, c)
+        prediction = self.Prediction.fc2(prediction)
+
+        return prediction

modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/__init__.py

@@ -0,0 +1 @@
+from .proxyless import plnas_linear_mix_se

modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/layers.py

@@ -0,0 +1,790 @@
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_same_padding(kernel_size):
+
+    if isinstance(kernel_size, tuple):
+        assert len(kernel_size) == 2, 'invalid kernel size: %s' % kernel_size
+        p1 = get_same_padding(kernel_size[0])
+        p2 = get_same_padding(kernel_size[1])
+        return p1, p2
+    assert isinstance(kernel_size,
+                      int), 'kernel size should be either `int` or `tuple`'
+    assert kernel_size % 2 > 0, 'kernel size should be odd number'
+    return kernel_size // 2
+
+
+def count_conv_flop(layer, x):
+    out_h = int(x.size(2) / layer.stride[0])
+    out_w = int(x.size(3) / layer.stride[1])
+    delta_ops = layer.in_channels * layer.out_channels * layer.kernel_size[
+        0] * layer.kernel_size[1] * out_h * out_w / layer.groups
+    return delta_ops
+
+
+def set_layer_from_config(layer_config):
+    if layer_config is None:
+        return None
+    name2layer = {
+        ZeroLayer.__name__: ZeroLayer,
+        MBInvertedConvLayer.__name__: MBInvertedConvLayer,
+        IdentityLayer.__name__: IdentityLayer,
+    }
+    layer_name = layer_config.pop('name')
+    layer = name2layer[layer_name]
+    return layer.build_from_config(layer_config)
+
+
+class MobileInvertedResidualBlock(nn.Module):
+
+    def __init__(self, mobile_inverted_conv, shortcut):
+        super(MobileInvertedResidualBlock, self).__init__()
+        self.mobile_inverted_conv = mobile_inverted_conv
+        self.shortcut = shortcut
+
+    def forward(self, x):
+        if self.mobile_inverted_conv.is_zero_layer():
+            res = x
+        elif self.shortcut is None or self.shortcut.is_zero_layer():
+            res = self.mobile_inverted_conv(x)
+        else:
+            conv_x = self.mobile_inverted_conv(x)
+            skip_x = self.shortcut(x)
+            res = skip_x + conv_x
+        return res
+
+    @property
+    def module_str(self):
+        return '(%s, %s)' % (self.mobile_inverted_conv.module_str,
+                             self.shortcut.module_str
+                             if self.shortcut is not None else None)
+
+    @property
+    def config(self):
+        return {
+            'name':
+            MobileInvertedResidualBlock.__name__,
+            'mobile_inverted_conv':
+            self.mobile_inverted_conv.config,
+            'shortcut':
+            self.shortcut.config if self.shortcut is not None else None,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        mobile_inverted_conv = set_layer_from_config(
+            config['mobile_inverted_conv'])
+        shortcut = set_layer_from_config(config['shortcut'])
+        return MobileInvertedResidualBlock(mobile_inverted_conv, shortcut)
+
+    def get_flops(self, x):
+        flops1, _ = self.mobile_inverted_conv.get_flops(x)
+        if self.shortcut:
+            flops2, _ = self.shortcut.get_flops(x)
+        else:
+            flops2 = 0
+        return flops1 + flops2, self.forward(x)
+
+
+class MBInvertedConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 stride=(1, 1),
+                 expand_ratio=6,
+                 mid_channels=None):
+        super(MBInvertedConvLayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+        if self.expand_ratio == 1:
+            self.inverted_bottleneck = None
+        else:
+            self.inverted_bottleneck = nn.Sequential(
+                OrderedDict([
+                    ('conv',
+                     nn.Conv2d(
+                         self.in_channels, feature_dim, 1, 1, 0, bias=False)),
+                    ('bn', nn.BatchNorm2d(feature_dim)),
+                    ('act', nn.PReLU()),
+                ]))
+        pad = get_same_padding(self.kernel_size)
+        self.depth_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(
+                     feature_dim,
+                     feature_dim,
+                     kernel_size,
+                     stride,
+                     pad,
+                     groups=feature_dim,
+                     bias=False)),
+                ('bn', nn.BatchNorm2d(feature_dim)),
+                ('act', nn.PReLU()),
+            ]))
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        if self.inverted_bottleneck:
+            x = self.inverted_bottleneck(x)
+        x = self.depth_conv(x)
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%dx%d_MBConv%d' % (self.kernel_size[0], self.kernel_size[1],
+                                   self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'kernel_size': self.kernel_size,
+            'stride': self.stride,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        if self.inverted_bottleneck:
+            total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+            x = self.inverted_bottleneck(x)
+        total_flops += count_conv_flop(self.depth_conv.conv, x)
+        x = self.depth_conv(x)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class IdentityLayer(nn.Module):
+
+    def __init__(self):
+        super(IdentityLayer, self).__init__()
+
+    def forward(self, x):
+        return x
+
+    @property
+    def module_str(self):
+        return 'Identity'
+
+    @property
+    def config(self):
+        return {
+            'name': IdentityLayer.__name__,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return IdentityLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        return 0, self.forward(x)
+
+
+class ZeroLayer(nn.Module):
+
+    def __init__(self, stride):
+        super(ZeroLayer, self).__init__()
+        self.stride = stride
+
+    def forward(self, x):
+        n, c, h, w = x.shape
+        h //= self.stride[0]
+        w //= self.stride[1]
+        device = x.device
+        padding = torch.zeros(n, c, h, w, device=device, requires_grad=False)
+        return padding
+
+    @property
+    def module_str(self):
+        return 'Zero'
+
+    @property
+    def config(self):
+        return {'name': ZeroLayer.__name__, 'stride': self.stride}
+
+    @staticmethod
+    def build_from_config(config):
+        return ZeroLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return True
+
+    def get_flops(self, x):
+        return 0, self.forward(x)
+
+
+def split_layer(total_channels, num_groups):
+    split = [
+        int(np.ceil(total_channels / num_groups)) for _ in range(num_groups)
+    ]
+    split[num_groups - 1] += total_channels - sum(split)
+    return split
+
+
+class MBInvertedMixConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 mix_conv_size=[1, 3, 5],
+                 stride=(1, 1),
+                 expand_ratio=6,
+                 mid_channels=None):
+        super(MBInvertedMixConvLayer, self).__init__()
+        self.mix_type = 0
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.mix_conv_size = mix_conv_size
+        self.stride = stride
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+        self.inverted_bottleneck = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(self.in_channels, feature_dim, 1, 1, 0,
+                           bias=False)),
+                ('bn', nn.BatchNorm2d(feature_dim)),
+                ('act', nn.PReLU()),
+            ]))
+
+        self.mix_conv_size = mix_conv_size
+        self.n_chunks = len(mix_conv_size)
+        self.split_in_channels = split_layer(feature_dim, self.n_chunks)
+
+        self.mix_conv = nn.ModuleList()
+        for idx in range(self.n_chunks):
+            kernel_size = self.mix_conv_size[idx]
+            pad = get_same_padding(kernel_size)
+            split_in_channels_ = self.split_in_channels[idx]
+            if self.mix_type == 0:
+                self.mix_conv.append(
+                    nn.Sequential(
+                        OrderedDict([
+                            ('conv',
+                             nn.Conv2d(
+                                 split_in_channels_,
+                                 split_in_channels_,
+                                 kernel_size,
+                                 stride,
+                                 pad,
+                                 groups=split_in_channels_,
+                                 bias=False)),
+                            ('bn', nn.BatchNorm2d(split_in_channels_)),
+                            ('act', nn.PReLU()),
+                        ])))
+            else:
+                self.mix_conv.append(
+                    nn.Sequential(
+                        OrderedDict([
+                            ('conv',
+                             nn.Conv2d(
+                                 split_in_channels_,
+                                 split_in_channels_,
+                                 kernel_size,
+                                 stride,
+                                 pad,
+                                 groups=split_in_channels_,
+                                 bias=False)),
+                            ('bn', nn.BatchNorm2d(split_in_channels_)),
+                        ])))
+        if self.mix_type != 0:
+            self.act = nn.PReLU()
+
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        x = self.inverted_bottleneck(x)
+        split = torch.split(x, self.split_in_channels, dim=1)
+        x = torch.cat([layer(s) for layer, s in zip(self.mix_conv, split)],
+                      dim=1)
+
+        if self.mix_type != 0:
+            x = self.act(x)
+        x = self.point_conv(x)
+
+        return x
+
+    @property
+    def module_str(self):
+        return '%s_MixConv%d' % (str(self.mix_conv_size), self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedMixConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'mix_conv_size': self.mix_conv_size,
+            'stride': self.stride,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedMixConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        # if self.inverted_bottleneck:
+        total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+        x = self.inverted_bottleneck(x)
+        split = torch.split(x, self.split_in_channels, dim=1)
+        out = []
+        # import pdb;pdb.set_trace()
+        for layer, s in zip(self.mix_conv, split):
+            out.append(layer(s))
+            total_flops += count_conv_flop(layer.conv, s)
+        x = torch.cat(out, dim=1)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class LinearMixConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 mix_conv_size=[1, 3, 5],
+                 stride=(1, 1),
+                 mid_channels=None):
+        super(LinearMixConvLayer, self).__init__()
+        self.mix_type = 0
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.mix_conv_size = mix_conv_size
+        self.stride = stride
+        self.mid_channels = mid_channels
+
+        self.mix_conv_size = mix_conv_size
+        self.n_chunks = len(mix_conv_size)
+        self.mix_conv = nn.ModuleList()
+
+        for idx in range(self.n_chunks):
+            kernel_size = self.mix_conv_size[idx]
+            pad = get_same_padding(kernel_size)
+            if self.mix_type == 0:
+                self.mix_conv.append(
+                    nn.Sequential(
+                        OrderedDict([
+                            ('conv',
+                             nn.Conv2d(
+                                 in_channels,
+                                 in_channels,
+                                 kernel_size,
+                                 stride,
+                                 pad,
+                                 groups=in_channels,
+                                 bias=False)),
+                            ('bn', nn.BatchNorm2d(in_channels)),
+                            ('act', nn.PReLU()),
+                        ])))
+            else:
+                self.mix_conv.append(
+                    nn.Sequential(
+                        OrderedDict([
+                            ('conv',
+                             nn.Conv2d(
+                                 in_channels,
+                                 in_channels,
+                                 kernel_size,
+                                 stride,
+                                 pad,
+                                 groups=in_channels,
+                                 bias=False)),
+                            ('bn', nn.BatchNorm2d(in_channels)),
+                        ])))
+
+        if self.mix_conv != 0:
+            self.act = nn.PReLU()
+
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(
+                     in_channels * self.n_chunks,
+                     out_channels,
+                     1,
+                     1,
+                     0,
+                     bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+    def forward(self, x):
+        x = torch.cat([layer(x) for layer in self.mix_conv], dim=1)
+        if self.mix_conv != 0:
+            x = self.act(x)
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%s_LinearMixConv' % (str(self.mix_conv_size))
+
+    @property
+    def config(self):
+        return {
+            'name': LinearMixConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'mix_conv_size': self.mix_conv_size,
+            'stride': self.stride,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return LinearMixConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        out = []
+        for layer in self.mix_conv:
+            out.append(layer(x))
+            total_flops += count_conv_flop(layer.conv, x)
+        x = torch.cat(out, dim=1)
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x
+
+
+class SELayer(nn.Module):
+    '''
+    '''
+
+    def __init__(self, input_channels: int, squeeze_factor: int = 4):
+        super(SELayer, self).__init__()
+        self.input_channels = input_channels
+        self.squeeze_factor = squeeze_factor
+        self.squeeze_channels = input_channels // squeeze_factor
+        self.fc1 = nn.Conv2d(
+            self.input_channels, self.squeeze_channels, 1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.fc2 = nn.Conv2d(
+            self.squeeze_channels, self.input_channels, 1, bias=True)
+
+    def _scale(self, input):
+        scale = F.adaptive_avg_pool2d(input, 1)
+        scale = self.fc1(scale)
+        scale = self.relu(scale)
+        scale = self.fc2(scale)
+        return torch.sigmoid(scale)
+
+    def forward(self, input):
+        scale = self._scale(input)
+        return scale * input
+
+    @property
+    def module_str(self):
+        return 'SE_%d' % (self.squeeze_factor)
+
+    @property
+    def config(self):
+        return {
+            'name': SELayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'squeeze_factor': self.squeeze_factor,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return SELayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def get_flops(self, x):
+        '''
+        count se flops, only compute the fc layers' calculation
+        '''
+        total_flops = 0
+        total_flops += self.input_channels * self.squeeze_channels * 2
+        b, c, h, w = x.shape
+        total_flops += c * h * w
+        return total_flops, x
+
+
+class MBInvertedRepConvLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 rep_conv_size=[1, 3, 5],
+                 stride=(1, 1),
+                 expand_ratio=6,
+                 mid_channels=None):
+        super(MBInvertedRepConvLayer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.rep_conv_size = rep_conv_size
+        self.stride = stride
+        self.expand_ratio = expand_ratio
+        self.mid_channels = mid_channels
+
+        feature_dim = round(
+            self.in_channels
+            * self.expand_ratio) if mid_channels is None else mid_channels
+
+        self.inverted_bottleneck = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(self.in_channels, feature_dim, 1, 1, 0,
+                           bias=False)),
+                ('bn', nn.BatchNorm2d(feature_dim)),
+                ('act', nn.PReLU()),
+            ]))
+
+        self.rep_conv_size = rep_conv_size
+        self.n_chunks = len(rep_conv_size)
+
+        self.rep_conv = nn.ModuleList()
+        for idx in range(self.n_chunks):
+            kernel_size = self.rep_conv_size[idx]
+            pad = get_same_padding(kernel_size)
+            self.rep_conv.append(
+                nn.Sequential(
+                    OrderedDict([
+                        ('conv',
+                         nn.Conv2d(
+                             feature_dim,
+                             feature_dim,
+                             kernel_size,
+                             stride,
+                             pad,
+                             groups=feature_dim,
+                             bias=False)),
+                        ('bn', nn.BatchNorm2d(feature_dim)),
+                    ])))
+
+        self.rep_conv_deploy = None
+        self.act = nn.PReLU()
+        self.point_conv = nn.Sequential(
+            OrderedDict([
+                ('conv',
+                 nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)),
+                ('bn', nn.BatchNorm2d(out_channels)),
+            ]))
+
+        self.deploy = False
+
+    def forward(self, x):
+        # if self.inverted_bottleneck:
+
+        x = self.inverted_bottleneck(x)
+
+        if not self.deploy:
+            out = []
+            for layer in self.rep_conv:
+                out.append(layer(x))
+            x = out[0]
+            for out_ in out[1:]:
+                x += out_
+        else:
+            x = self.rep_conv_deploy(x)
+        x = self.act(x)
+        x = self.point_conv(x)
+        return x
+
+    @property
+    def module_str(self):
+        return '%s_RepConv%d' % (str(self.rep_conv_size), self.expand_ratio)
+
+    @property
+    def config(self):
+        return {
+            'name': MBInvertedMixConvLayer.__name__,
+            'in_channels': self.in_channels,
+            'out_channels': self.out_channels,
+            'rep_conv_size': self.rep_conv_size,
+            'stride': self.stride,
+            'expand_ratio': self.expand_ratio,
+            'mid_channels': self.mid_channels,
+        }
+
+    @staticmethod
+    def build_from_config(config):
+        return MBInvertedMixConvLayer(**config)
+
+    @staticmethod
+    def is_zero_layer():
+        return False
+
+    def switch_to_deploy(self):
+        self.deploy = True
+        feature_dim = self.rep_conv[0].conv.in_channels
+        stride = self.rep_conv[0].conv.stride
+        kernel_size = max(self.rep_conv_size)
+        pad = get_same_padding(kernel_size)
+        self.rep_conv_deploy = nn.Conv2d(
+            feature_dim,
+            feature_dim,
+            kernel_size,
+            stride,
+            pad,
+            groups=feature_dim,
+            bias=True)
+
+        kernel, bias = self.get_equivalent_kernel_bias()
+        self.rep_conv_deploy.weight.data = kernel
+        self.rep_conv_deploy.bias.data = bias
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__('rep_conv')
+
+    def get_equivalent_kernel_bias(self):
+        import pdb
+        pdb.set_trace()
+        if 1 in self.rep_conv_size:
+            kernel1x1, bias1x1 = self._fuse_bn_tensor(
+                self.rep_conv[self.rep_conv_size.index(1)])
+        else:
+            kernel1x1 = None
+            bias1x1 = 0
+
+        if 3 in self.rep_conv_size:
+            kernel3x3, bias3x3 = self._fuse_bn_tensor(
+                self.rep_conv[self.rep_conv_size.index(3)])
+        else:
+            kernel3x3 = None
+            bias3x3 = 0
+
+        if 5 in self.rep_conv_size:
+            kernel5x5, bias5x5 = self._fuse_bn_tensor(
+                self.rep_conv[self.rep_conv_size.index(5)])
+
+        else:
+            kernel5x5 = 0
+            bias5x5 = 0
+
+        return kernel5x5 + self._pad_1x1_to_5x5_tensor(
+            kernel1x1) + self._pad_3x3_to_5x5_tensor(
+                kernel3x3), bias5x5 + bias3x3 + bias1x1
+
+    def _pad_1x1_to_5x5_tensor(self, kernel1x1):
+        if kernel1x1 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel1x1, [2, 2, 2, 2])
+
+    def _pad_3x3_to_5x5_tensor(self, kernel3x3):
+        if kernel3x3 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel3x3, [1, 1, 1, 1])
+
+    def _fuse_bn_tensor(self, branch):
+        if branch is None:
+            return 0, 0
+        if isinstance(branch, nn.Sequential):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, 'id_tensor'):
+                input_dim = self.in_channels // self.groups
+                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3),
+                                        dtype=np.float32)
+                for i in range(self.in_channels):
+                    kernel_value[i, i % input_dim, 1, 1] = 1
+                self.id_tensor = torch.from_numpy(kernel_value).to(
+                    branch.weight.device)
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+
+        return kernel * t, beta - running_mean * gamma / std
+
+    def get_flops(self, x):
+        '''count conv flops, skip BN and other small flops
+        '''
+        total_flops = 0
+        total_flops += count_conv_flop(self.inverted_bottleneck.conv, x)
+        x = self.inverted_bottleneck(x)
+
+        total_flops += count_conv_flop(self.rep_conv[-1].conv, x)
+        out = []
+        for layer in self.rep_conv:
+            out.append(layer(x))
+        x = out[0]
+        for out_ in out[1:]:
+            x += out_
+        total_flops += count_conv_flop(self.point_conv.conv, x)
+        x = self.point_conv(x)
+        return total_flops, x

modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/mix_ops.py

@@ -0,0 +1,226 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+
+from .layers import (LinearMixConvLayer, MBInvertedConvLayer,
+                     MBInvertedMixConvLayer, MBInvertedRepConvLayer, SELayer,
+                     ZeroLayer)
+
+
+def detach_variable(inputs):
+    if isinstance(inputs, tuple):
+        return tuple([detach_variable(x) for x in inputs])
+    else:
+        x = inputs.detach()
+        x.requires_grad = inputs.requires_grad
+        return x
+
+
+def delta_ij(i, j):
+    if i == j:
+        return 1
+    else:
+        return 0
+
+
+def conv_func_by_name(name):
+
+    name2ops = {
+        'Identity':
+        lambda in_C, out_C, S, height: IdentityLayer(
+            in_C, out_C, ops_order=ops_order),
+        'Zero':
+        lambda in_C, out_C, S, height: ZeroLayer(stride=S),
+    }
+    name2ops.update({
+        '3x3_MBConv1':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 3), 3), S, 1),
+        '3x3_MBConv2':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 3), 3), S, 2),
+        '3x3_MBConv3':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 3), 3), S, 3),
+        '3x3_MBConv4':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 3), 3), S, 4),
+        '3x3_MBConv5':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 3), 3), S, 5),
+        '3x3_MBConv6':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 3), 3), S, 6),
+        #######################################################################################
+        '5x5_MBConv1':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 5), 5), S, 1),
+        '5x5_MBConv2':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 5), 5), S, 2),
+        '5x5_MBConv3':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 5), 5), S, 3),
+        '5x5_MBConv4':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 5), 5), S, 4),
+        '5x5_MBConv5':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 5), 5), S, 5),
+        '5x5_MBConv6':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 5), 5), S, 6),
+        #######################################################################################
+        '7x7_MBConv1':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 7), 7), S, 1),
+        '7x7_MBConv2':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 7), 7), S, 2),
+        '7x7_MBConv3':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 7), 7), S, 3),
+        '7x7_MBConv4':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 7), 7), S, 4),
+        '7x7_MBConv5':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 7), 7), S, 5),
+        '7x7_MBConv6':
+        lambda in_C, out_C, S, height: MBInvertedConvLayer(
+            in_C, out_C, (min(height, 7), 7), S, 6),
+        #######################################################################################
+        '13_MixConv1':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 1),
+        '13_MixConv2':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 2),
+        '13_MixConv3':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 3),
+        '13_MixConv4':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 4),
+        '13_MixConv5':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 5),
+        '13_MixConv6':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 6),
+        #######################################################################################
+        '35_MixConv1':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 1),
+        '35_MixConv2':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 2),
+        '35_MixConv3':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 3),
+        '35_MixConv4':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 4),
+        '35_MixConv5':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 5),
+        '35_MixConv6':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 6),
+        #######################################################################################
+        '135_MixConv1':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 1),
+        '135_MixConv2':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 2),
+        '135_MixConv3':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 3),
+        '135_MixConv4':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 4),
+        '135_MixConv5':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 5),
+        '135_MixConv6':
+        lambda in_C, out_C, S, height: MBInvertedMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 6),
+        #######################################################################################
+        '13_LinMixConv':
+        lambda in_C, out_C, S, height: LinearMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S),
+        '35_LinMixConv':
+        lambda in_C, out_C, S, height: LinearMixConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S),
+        '135_LinMixConv':
+        lambda in_C, out_C, S, height: LinearMixConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S),
+        #######################################################################################
+        'SE_2':
+        lambda in_C, out_C, S, height: SELayer(in_C, 2),
+        'SE_4':
+        lambda in_C, out_C, S, height: SELayer(in_C, 4),
+        'SE_8':
+        lambda in_C, out_C, S, height: SELayer(in_C, 8),
+        #######################################################################################
+        '13_RepConv1':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 1),
+        '13_RepConv2':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 2),
+        '13_RepConv3':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 3),
+        '13_RepConv4':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 4),
+        '13_RepConv5':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 5),
+        '13_RepConv6':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3)], S, 6),
+        #######################################################################################
+        '35_RepConv1':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 1),
+        '35_RepConv2':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 2),
+        '35_RepConv3':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 3),
+        '35_RepConv4':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 4),
+        '35_RepConv5':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 5),
+        '35_RepConv6':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [(min(height, 3), 3), (min(height, 5), 5)], S, 6),
+        #######################################################################################
+        '135_RepConv1':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 1),
+        '135_RepConv2':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 2),
+        '135_RepConv3':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 3),
+        '135_RepConv4':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 4),
+        '135_RepConv5':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 5),
+        '135_RepConv6':
+        lambda in_C, out_C, S, height: MBInvertedRepConvLayer(
+            in_C, out_C, [1, (min(height, 3), 3), (min(height, 5), 5)], S, 6),
+    })
+    return name2ops[name]

modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py

@@ -0,0 +1,138 @@
+# Part of the implementation is borrowed and modified from ProxylessNAS,
+# publicly available at https://github.com/mit-han-lab/proxylessnas
+# paper linking at https://arxiv.org/pdf/1812.00332.pdf
+import re
+import sys
+from queue import Queue
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .layers import IdentityLayer, MobileInvertedResidualBlock
+from .mix_ops import conv_func_by_name
+
+
+class NasRecBackbone(nn.Module):
+
+    def __init__(self, first_conv, blocks):
+        super(NasRecBackbone, self).__init__()
+        self.first_conv = first_conv
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, x):
+        x = self.first_conv(x)
+        for block in self.blocks:
+            x = block(x)
+        return x
+
+    def get_flops(self, x):
+        expected_flops = 0
+        # first conv
+        flop = count_conv_flop(self.first_conv[0], x)
+        x = self.first_conv(x)
+        expected_flops += flop
+        # blocks
+        for mb_conv in self.blocks:
+            assert isinstance(mb_conv, MobileInvertedResidualBlock)
+            if mb_conv.shortcut is None:
+                shortcut_flop = 0
+            else:
+                shortcut_flop, _ = mb_conv.shortcut.get_flops(x)
+            expected_flops += shortcut_flop
+            expected_flops += mb_conv.get_flops(x)[0]
+
+            x = mb_conv(x)
+        return expected_flops
+
+
+class CompactRecBackboneMixSE(NasRecBackbone):
+
+    def __init__(self, first_stride, input_channel, stride_stages,
+                 n_cell_stages, width_stages, conv_op_ids, conv_candidates,
+                 se_candidates):
+        input_block_channel = 24
+        first_conv = nn.Sequential(
+            nn.Conv2d(
+                input_channel,
+                input_block_channel,
+                kernel_size=(3, 3),
+                stride=first_stride,
+                padding=1,
+                bias=False), nn.BatchNorm2d(input_block_channel), nn.PReLU())
+
+        assert len(conv_op_ids) - 4 == sum(n_cell_stages)
+        blocks = []
+        img_height = 16
+        height_flag = 0
+        for width, n_cell, s in zip(width_stages, n_cell_stages,
+                                    stride_stages):
+            for i in range(n_cell):
+                if i == 1:
+                    img_height = int(img_height / 2)
+
+                if img_height % 2 == 0:
+                    height_flag = 1
+                else:
+                    height_flag = 0
+
+                if i == 0:
+                    stride = s
+                else:
+                    stride = (1, 1)
+                block_i = len(blocks)
+                conv_op = conv_func_by_name(
+                    conv_candidates[conv_op_ids[block_i]])(
+                        input_block_channel, width, stride,
+                        img_height + height_flag)
+                if stride == (1, 1) and input_block_channel == width:
+                    shortcut = IdentityLayer()
+                else:
+                    shortcut = None
+                inverted_residual_block = MobileInvertedResidualBlock(
+                    conv_op, shortcut)
+                blocks.append(inverted_residual_block)
+                input_block_channel = width
+
+            block_i = len(blocks)
+            se_op = conv_func_by_name(se_candidates[conv_op_ids[block_i]])(
+                input_block_channel, width, stride, img_height)
+
+            inverted_residual_block = MobileInvertedResidualBlock(se_op, None)
+            blocks.append(inverted_residual_block)
+        self.out_channel = input_block_channel
+
+        super(CompactRecBackboneMixSE, self).__init__(first_conv, blocks)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+
+def plnas_linear_mix_se(input_channel, output_channel):
+    conv_candidates = [
+        '5x5_MBConv2', '5x5_MBConv4', '5x5_MBConv6', '3x3_MBConv2',
+        '3x3_MBConv4', '3x3_MBConv6', '13_MixConv2', '13_MixConv4',
+        '13_MixConv6', '35_MixConv2', '35_MixConv4', '35_MixConv6',
+        '135_MixConv2', '135_MixConv4', '135_MixConv6', '13_LinMixConv',
+        '35_LinMixConv', '135_LinMixConv', '13_RepConv2', '13_RepConv4',
+        '13_RepConv6', '35_RepConv2', '35_RepConv4', '35_RepConv6',
+        '135_RepConv2', '135_RepConv4', '135_RepConv6', 'Zero'
+    ]
+    se_candidates = ['SE_2', 'SE_4', 'SE_8', 'Zero']
+
+    stride_stages = [(2, 2), (2, 1), (2, 1), (2, 1)]
+    n_cell_stages = [5, 5, 5, 5]
+    width_stages = [32, 64, 96, 123]
+    conv_op_ids = [
+        2, 23, 24, 26, 2, 2, 11, 27, 27, 27, 27, 2, 0, 2, 16, 10, 27, 2, 2, 2,
+        22, 10, 27, 3
+    ]
+    net = CompactRecBackboneMixSE(2, input_channel, stride_stages,
+                                  n_cell_stages, width_stages, conv_op_ids,
+                                  conv_candidates, se_candidates)
+
+    return net

modelscope/models/cv/ocr_recognition/preprocessor.py

@@ -31,8 +31,6 @@ class OCRRecognitionPreprocessor(Preprocessor):
         self.do_chunking = cfgs.model.inference_kwargs.do_chunking
         self.target_height = cfgs.model.inference_kwargs.img_height
         self.target_width = cfgs.model.inference_kwargs.img_width
-        if self.do_chunking:
-            self.target_width = 804
 
     def keepratio_resize(self, img):
         cur_ratio = img.shape[1] / float(img.shape[0])
@@ -45,8 +43,8 @@ class OCRRecognitionPreprocessor(Preprocessor):
             cur_target_height = self.target_height
             cur_target_width = int(self.target_height * cur_ratio)
         img = cv2.resize(img, (cur_target_width, cur_target_height))
-        mask = np.zeros([mask_height, mask_width]).astype(np.uint8)
-        mask[:img.shape[0], :img.shape[1]] = img
+        mask = np.zeros([mask_height, mask_width, 3]).astype(np.uint8)
+        mask[:img.shape[0], :img.shape[1], :] = img
         img = mask
         return img
 
@@ -65,29 +63,32 @@ class OCRRecognitionPreprocessor(Preprocessor):
         data_batch = []
         for item in inputs:
             if isinstance(item, str):
-                img = np.array(load_image(item).convert('L'))
+                img = np.array(load_image(item).convert('RGB'))
             elif isinstance(item, PIL.Image.Image):
-                img = np.array(item.convert('L'))
+                img = np.array(item.convert('RGB'))
             elif isinstance(item, np.ndarray):
-                if len(item.shape) == 3:
-                    img = cv2.cvtColor(item, cv2.COLOR_RGB2GRAY)
+                if len(item.shape) == 2:
+                    img = cv2.cvtColor(item, cv2.COLOR_GRAY2RGB)
+                else:
+                    img = item
             else:
                 raise TypeError(
                     f'inputs should be either (a list of) str, PIL.Image, np.array, but got {type(item)}'
                 )
-
             img = self.keepratio_resize(img)
             img = torch.FloatTensor(img)
             if self.do_chunking:
                 chunk_img = []
                 for i in range(3):
                     left = (300 - 48) * i
-                    chunk_img.append(img[:, left:left + 300])
+                    chunk_img.append(img[:, left:left + 300, :])
                 merge_img = torch.cat(chunk_img, 0)
-                data = merge_img.view(3, 1, self.target_height, 300) / 255.
+                data = merge_img.view(3, self.target_height, 300, 3) / 255.
+                data = data.permute(0, 3, 1, 2)
             else:
-                data = img.view(1, 1, self.target_height,
-                                self.target_width) / 255.
+                data = img.view(1, self.target_height, self.target_width,
+                                3) / 255.
+                data = data.permute(0, 3, 1, 2)
             data_batch.append(data)
         data_batch = torch.cat(data_batch, 0)
-        return data_batch
+        return {'image': data_batch}

modelscope/models/cv/table_recognition/lineless_table_process.py

@@ -346,6 +346,7 @@ def merge_outputs(detections):
 def filter(results, logi, ps):
     # this function select boxes
     batch_size, feat_dim = logi.shape[0], logi.shape[2]
+
     num_valid = sum(results[1][:, 8] >= 0.15)
 
     slct_logi = np.zeros((batch_size, num_valid, feat_dim), dtype=np.float32)
@@ -358,6 +359,14 @@ def filter(results, logi, ps):
     return torch.Tensor(slct_logi).cuda(), torch.Tensor(slct_dets).cuda()
 
 
+def normalized_ps(ps, vocab_size):
+    ps = torch.round(ps).to(torch.int64)
+    ps = torch.where(ps < vocab_size, ps, (vocab_size - 1)
+                     * torch.ones(ps.shape).to(torch.int64).cuda())
+    ps = torch.where(ps >= 0, ps, torch.zeros(ps.shape).to(torch.int64).cuda())
+    return ps
+
+
 def process_detect_output(output, meta):
     K, MK = 3000, 5000
     num_classes = 2
@@ -390,6 +399,7 @@ def process_detect_output(output, meta):
     logi = logi + cr
     results = merge_outputs(detections)
     slct_logi_feat, slct_dets_feat = filter(results, logi, raw_dets[:, :, :8])
+    slct_dets_feat = normalized_ps(slct_dets_feat, 256)
     slct_output_dets = results[1][:slct_logi_feat.shape[1], :8]
 
     return slct_logi_feat, slct_dets_feat, slct_output_dets

modelscope/models/multi_modal/stable_diffusion/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .stable_diffusion import StableDiffusion

modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py

@@ -0,0 +1,152 @@
+# Copyright 2023-2024 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+import os
+from functools import partial
+from typing import Callable, List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from modelscope.metainfo import Models
+from modelscope.models import TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.outputs import OutputKeys
+from modelscope.utils.checkpoint import save_checkpoint, save_configuration
+from modelscope.utils.constant import Tasks
+
+
+@MODELS.register_module(
+    Tasks.text_to_image_synthesis, module_name=Models.stable_diffusion)
+class StableDiffusion(TorchModel):
+    """ The implementation of efficient diffusion tuning model based on TorchModel.
+
+    This model is constructed with the implementation of stable diffusion model. If you want to
+    finetune lightweight parameters on your own dataset, you can define you own tuner module
+    and load in this cls.
+    """
+
+    def __init__(self, model_dir, *args, **kwargs):
+        """ Initialize a vision efficient diffusion tuning model.
+
+        Args:
+          model_dir: model id or path, where model_dir/pytorch_model.bin
+        """
+        super().__init__(model_dir, *args, **kwargs)
+        pretrained_model_name_or_path = kwargs.pop(
+            'pretrained_model_name_or_path', 'runwayml/stable-diffusion-v1-5')
+        revision = kwargs.pop('revision', None)
+        self.lora_tune = kwargs.pop('lora_tune', True)
+
+        self.weight_dtype = torch.float32
+        self.device = torch.device(
+            'cuda' if torch.cuda.is_available() else 'cpu')
+
+        # Load scheduler, tokenizer and models
+        self.noise_scheduler = DDPMScheduler.from_pretrained(
+            pretrained_model_name_or_path, subfolder='scheduler')
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder='tokenizer',
+            revision=revision)
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder='text_encoder',
+            revision=revision)
+        self.vae = AutoencoderKL.from_pretrained(
+            pretrained_model_name_or_path, subfolder='vae', revision=revision)
+        self.unet = UNet2DConditionModel.from_pretrained(
+            pretrained_model_name_or_path, subfolder='unet', revision=revision)
+
+        # Freeze gradient calculation and move to device
+        if self.vae is not None:
+            self.vae.requires_grad_(False)
+            self.vae = self.vae.to(self.device)
+        if self.text_encoder is not None:
+            self.text_encoder.requires_grad_(False)
+            self.text_encoder = self.text_encoder.to(self.device)
+        if self.unet is not None:
+            if self.lora_tune:
+                self.unet.requires_grad_(False)
+            self.unet = self.unet.to(self.device)
+
+    def tokenize_caption(self, captions):
+        """ Convert caption text to token data.
+
+        Args:
+          captions: a batch of texts.
+        Returns: token's data as tensor.
+        """
+        inputs = self.tokenizer(
+            captions,
+            max_length=self.tokenizer.model_max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt')
+        return inputs.input_ids
+
+    def forward(self, text='', target=None):
+        self.unet.train()
+        self.unet = self.unet.to(self.device)
+
+        with torch.no_grad():
+            latents = self.vae.encode(
+                target.to(dtype=self.weight_dtype)).latent_dist.sample()
+        latents = latents * self.vae.config.scaling_factor
+
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(latents)
+        bsz = latents.shape[0]
+        # Sample a random timestep for each image
+        timesteps = torch.randint(
+            0,
+            self.noise_scheduler.num_train_timesteps, (bsz, ),
+            device=latents.device)
+        timesteps = timesteps.long()
+
+        # Add noise to the latents according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+        noisy_latents = self.noise_scheduler.add_noise(latents, noise,
+                                                       timesteps)
+
+        input_ids = self.tokenize_caption(text).to(self.device)
+
+        # Get the text embedding for conditioning
+        with torch.no_grad():
+            encoder_hidden_states = self.text_encoder(input_ids)[0]
+
+        # Get the target for loss depending on the prediction type
+        if self.noise_scheduler.config.prediction_type == 'epsilon':
+            target = noise
+        elif self.noise_scheduler.config.prediction_type == 'v_prediction':
+            target = self.noise_scheduler.get_velocity(latents, noise,
+                                                       timesteps)
+        else:
+            raise ValueError(
+                f'Unknown prediction type {self.noise_scheduler.config.prediction_type}'
+            )
+
+        # Predict the noise residual and compute loss
+        model_pred = self.unet(noisy_latents, timesteps,
+                               encoder_hidden_states).sample
+
+        loss = F.mse_loss(model_pred.float(), target.float(), reduction='mean')
+
+        output = {OutputKeys.LOSS: loss}
+        return output
+
+    def save_pretrained(self,
+                        target_folder: Union[str, os.PathLike],
+                        save_checkpoint_names: Union[str, List[str]] = None,
+                        save_function: Callable = partial(
+                            save_checkpoint, with_meta=False),
+                        config: Optional[dict] = None,
+                        save_config_function: Callable = save_configuration,
+                        **kwargs):
+        # Save only the lora model, skip saving and copying the original weights
+        if self.lora_tune:
+            pass
+        else:
+            super().save_pretrained(target_folder, save_checkpoint_names,
+                                    save_function, config,
+                                    save_config_function, **kwargs)

modelscope/models/nlp/gpt3/backbone.py

@@ -353,7 +353,7 @@ class GPT3Model(PreTrainedModel):
         model.load_state_dict(state_dict)
         return model
 
-    def generate(self, tokens, temperature=1.0, **kwargs):
+    def streaming_generate(self, tokens, temperature=1.0, **kwargs):
         top_k = kwargs.pop('top_k', self.config.top_k)
         top_p = kwargs.pop('top_p', self.config.top_p)
         max_length = kwargs.pop('max_length', tokens.size(1) + 100)
@@ -410,6 +410,9 @@ class GPT3Model(PreTrainedModel):
                 # Update the tokens.
                 tokens[started, context_length] = new_sample[started]
 
+                yield TokenGeneratorOutput(sequences=tokens[:, :(context_length
+                                                                 + 1)])
+
                 done_token = (new_sample == termination_id).byte() & \
                     started.byte()
 
@@ -419,5 +422,8 @@ class GPT3Model(PreTrainedModel):
                 if done:
                     break
 
-        tokens = tokens[:, :(context_length + 1)]
-        return TokenGeneratorOutput(sequences=tokens)
+    def generate(self, tokens, temperature=1.0, **kwargs):
+        last_output = None
+        for output in self.streaming_generate(tokens, temperature, **kwargs):
+            last_output = output
+        return last_output

modelscope/models/nlp/gpt3/distributed_gpt3.py

@@ -33,6 +33,7 @@ from modelscope.models.nlp.gpt3 import GPT3Config
 from modelscope.outputs import TextGenerationModelOutput, TokenGeneratorOutput
 from modelscope.utils.megatron_utils import init_megatron_util
 from modelscope.utils.nlp.load_checkpoint import pre_load
+from modelscope.utils.streaming_output import StreamingOutputMixin
 
 
 class GPT3ParallelMLP(nn.Module):
@@ -945,7 +946,7 @@ def split_state_dict(state_dict: Dict[str, torch.Tensor], model: GPT3Model,
     return state_dict
 
 
-class DistributedGPT3(TorchModel):
+class DistributedGPT3(TorchModel, StreamingOutputMixin):
 
     def __init__(self,
                  model_dir,
@@ -1022,7 +1023,11 @@ class DistributedGPT3(TorchModel):
 
             losses = losses.float()
             loss_mask = loss_mask.view(-1).float()
-            loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
+            mask_sum = loss_mask.sum()
+            if mask_sum == 0:
+                loss = torch.sum(losses.view(-1)).zero_()
+            else:
+                loss = torch.sum(losses.view(-1) * loss_mask) / mask_sum
 
         return TextGenerationModelOutput(logits=logits, loss=loss)
 
@@ -1104,6 +1109,10 @@ class DistributedGPT3(TorchModel):
             # Update the tokens.
             tokens[started, context_length] = new_sample[started]
 
+            # streaming output
+            yield TokenGeneratorOutput(sequences=tokens[:, :(context_length
+                                                             + 1)])
+
             # Update the context length for the next token generation.
             prev_context_length = context_length
 
@@ -1128,9 +1137,6 @@ class DistributedGPT3(TorchModel):
             if use_eod_token_for_early_termination and done:
                 break
 
-        tokens = tokens[:, :(context_length + 1)]
-        return TokenGeneratorOutput(sequences=tokens)
-
     def beam_search(self, tokens, beam_size=5, num_return_gen=1, **kwargs):
         batch_size = tokens.size(0)
         assert (batch_size == 1)
@@ -1247,10 +1253,17 @@ class DistributedGPT3(TorchModel):
     @torch.no_grad()
     def generate(self, tokens, do_sample=True, *args, **kwargs):
         if do_sample:
-            return self.sample(tokens, *args, **kwargs)
+            last_output = None
+            for output in self.sample(tokens, *args, **kwargs):
+                last_output = output
+            return last_output
         else:
             return self.beam_search(tokens, *args, **kwargs)
 
+    @torch.no_grad()
+    def stream(self, tokens, *args, **kwargs):
+        return self.sample(tokens, *args, **kwargs)
+
     def state_dict(self, destination=None, prefix='', keep_vars=False):
         return self.dist_model.state_dict(destination, prefix, keep_vars)
 

modelscope/models/nlp/gpt3/text_generation.py

@@ -1,6 +1,6 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 from collections import OrderedDict
-from typing import Dict
+from typing import Dict, Generator
 
 import torch
 from transformers import BertTokenizer
@@ -11,12 +11,13 @@ from modelscope.models.builder import MODELS
 from modelscope.models.nlp.gpt3 import GPT3Model
 from modelscope.utils.constant import Tasks
 from modelscope.utils.hub import read_config
+from modelscope.utils.streaming_output import StreamingOutputMixin
 
 __all__ = ['GPT3ForTextGeneration']
 
 
 @MODELS.register_module(Tasks.text_generation, module_name=Models.gpt3)
-class GPT3ForTextGeneration(TorchModel):
+class GPT3ForTextGeneration(TorchModel, StreamingOutputMixin):
 
     def __init__(self, model_dir: str, *args, **kwargs):
         """initialize the text generation model from the `model_dir` path.
@@ -77,3 +78,9 @@ class GPT3ForTextGeneration(TorchModel):
                         state_dict: 'OrderedDict[str, Tensor]',
                         strict: bool = True):
         return self.model.load_state_dict(state_dict, strict)
+
+    def stream(self, inputs, **kwargs) -> Generator:
+        tokens = inputs['input_ids']
+        lengths = self._get_length(inputs['attention_mask'])
+        return self.model.streaming_generate(
+            tokens, prompt_length=lengths, **kwargs)

modelscope/msdatasets/data_files/data_files_manager.py

@@ -58,6 +58,7 @@ class DataFilesManager(object):
 
         # Set context. Note: no need to update context_config.
         download_config.oss_config = self.oss_config
+        download_config.num_proc = self.input_config_kwargs.get('num_proc', 4)
         dataset_context_config.download_config = download_config
         self.dataset_context_config = dataset_context_config
         os.makedirs(download_config.cache_dir, exist_ok=True)

modelscope/msdatasets/data_loader/data_loader.py

@@ -86,7 +86,6 @@ class OssDownloader(BaseDownloader):
     def _authorize(self) -> None:
         """ Authorization of target dataset.
         Get credentials from cache and send to the modelscope-hub in the future. """
-        # TODO: obtain credentials from loacl cache when available.
         cookies = ModelScopeConfig.get_cookies()
         git_token = ModelScopeConfig.get_token()
         user_info = ModelScopeConfig.get_user_info()

modelscope/msdatasets/data_loader/data_loader_manager.py

@@ -127,15 +127,15 @@ class RemoteDataLoaderManager(DataLoaderManager):
             return dataset_ret
         # To use the modelscope data loader
         elif data_loader_type == RemoteDataLoaderType.MS_DATA_LOADER:
-            oss_data_loader = OssDownloader(
+            oss_downloader = OssDownloader(
                 dataset_context_config=self.dataset_context_config)
-            oss_data_loader.process()
+            oss_downloader.process()
             # download statistics
             self.api.dataset_download_statistics(
                 dataset_name=dataset_name,
                 namespace=namespace,
                 use_streaming=use_streaming)
-            return oss_data_loader.dataset
+            return oss_downloader.dataset
         else:
             raise f'Expected remote data loader type: {RemoteDataLoaderType.HF_DATA_LOADER.value}/' \
                   f'{RemoteDataLoaderType.MS_DATA_LOADER.value}, but got {data_loader_type} .'

modelscope/msdatasets/dataset_cls/custom_datasets/ocr_recognition_dataset.py

@@ -68,7 +68,7 @@ class OCRRecognitionDataset(TorchCustomDataset):
             buf.seek(0)
             img = Image.open(buf).convert('L')
             if self.reco_preprocess is not None:
-                img = self.reco_preprocess(img)
+                img = self.reco_preprocess(img)['image']
 
             label_key = 'label-%09d' % index
             label = txn.get(label_key.encode()).decode('utf-8')

modelscope/msdatasets/dataset_cls/dataset.py

@@ -6,6 +6,7 @@ import os
 import datasets
 import pandas as pd
 from datasets import IterableDataset
+from tqdm import tqdm
 
 from modelscope.msdatasets.utils.maxcompute_utils import MaxComputeUtil
 from modelscope.utils.constant import (DEFAULT_MAXCOMPUTE_ENDPOINT,
@@ -23,15 +24,18 @@ class ExternalDataset(object):
     def __init__(self, split_path_dict, config_kwargs):
         self.split_path_dict = split_path_dict
         self.config_kwargs = copy.deepcopy(config_kwargs)
-        self.config_kwargs.update({'split_config': split_path_dict})
+        self.config_kwargs.update({'split_config': self.split_path_dict})
         # dataset for specific extensions
         self.spec_extension_dataset = None
-        self.split_data_files = {k: [] for k, _ in split_path_dict.items()}
+        self.split_data_files = {
+            k: []
+            for k, _ in self.split_path_dict.items()
+        }
         self.custom_map = {}
 
         # the extension of file
         file_ext = ''
-        for split_name, split_dir in split_path_dict.items():
+        for split_name, split_dir in self.split_path_dict.items():
             if isinstance(split_dir, str) and os.path.isdir(split_dir):
                 split_file_names = os.listdir(split_dir)
                 set_files_exts = set([
@@ -91,28 +95,52 @@ class NativeIterableDataset(IterableDataset):
         super().__init__(ex_iterable=ex_iterable, info=info, split=split)
 
     def __iter__(self):
-        for key, entity in self._iter():
+        for key, entity in tqdm(
+                self._iter(),
+                desc='Overall progress',
+                total=self.n_shards,
+                dynamic_ncols=True):
             if isinstance(entity, dict):
                 ret = {}
-                for k, v in entity.items():
-                    ret[k] = v
-                    if k.endswith(':FILE'):
-                        dl_manager = self._ex_iterable.kwargs.get('dl_manager')
-                        ex_cache_path = dl_manager.download_and_extract(v)
-                        ret[k] = ex_cache_path
-                        if k.endswith('Image:FILE'):
-                            from PIL import Image
-                            ret[k + ':Object'] = Image.open(fp=ex_cache_path)
-                        if k.endswith('Audio:FILE'):
-                            import torchaudio
-                            waveform_and_rate = torchaudio.load(ex_cache_path)
-                            ret[k + ':Object'] = waveform_and_rate
+                try:
+                    for k, v in entity.items():
+                        ret[k] = v
+                        if k.endswith(':FILE'):
+                            dl_manager = self._ex_iterable.kwargs.get(
+                                'dl_manager')
+                            ex_cache_path = dl_manager.download_and_extract(v)
+                            ret[k] = ex_cache_path
+                            if k.endswith('Image:FILE'):
+                                from PIL import Image
+                                ret[k
+                                    + ':Object'] = Image.open(fp=ex_cache_path)
+                            if k.endswith('Audio:FILE'):
+                                import torchaudio
+                                waveform_and_rate = torchaudio.load(
+                                    ex_cache_path)
+                                ret[k + ':Object'] = waveform_and_rate
+                except Exception as e:
+                    logger.error(e)
+                    ret = {}
+
                 entity = ret
 
             yield entity
 
     def __len__(self):
-        return 1
+        return self.n_shards
+
+    def head(self, n=5):
+        """
+        Returns the first n rows of the dataset.
+
+        Args:
+            n (int): Number of rows to return.
+
+        Returns:
+            Dict[str, list]: e.g. {'col1': [val11, val12, ...], 'col2': [val21, val22, ...]}
+        """
+        return self._head(n=n)
 
 
 class VirgoDataset(object):

modelscope/msdatasets/download/dataset_builder.py

@@ -6,8 +6,9 @@ from typing import Dict, Union
 import datasets
 import pandas as pd
 import pyarrow as pa
-from datasets import (ArrowBasedBuilder, GeneratorBasedBuilder,
-                      IterableDataset, IterableDatasetDict)
+from datasets import (ArrowBasedBuilder, Dataset, DatasetDict,
+                      GeneratorBasedBuilder, IterableDataset,
+                      IterableDatasetDict)
 from datasets.filesystems import is_remote_filesystem
 from datasets.info import DatasetInfo
 from datasets.naming import camelcase_to_snakecase
@@ -47,6 +48,7 @@ class CsvDatasetBuilder(csv.Csv):
         self.meta_data_files = dataset_context_config.data_meta_config.meta_data_files
         self.zip_data_files = dataset_context_config.data_meta_config.zip_data_files
         self.input_config_kwargs = dataset_context_config.config_kwargs
+        self.split_path_dict = dict({})
 
         self.cache_build_dir = os.path.join(self.cache_root_dir,
                                             self.namespace, self.dataset_name,
@@ -61,16 +63,25 @@ class CsvDatasetBuilder(csv.Csv):
         sub_dir_hash = get_subdir_hash_from_split(
             split=split, version=self.version)
 
+        from datasets.data_files import DataFilesDict, DataFilesList
+        data_files = {
+            k: DataFilesList([v], origin_metadata=None)
+            for k, v in self.meta_data_files.items()
+        }
+        data_files = DataFilesDict.from_local_or_remote(data_files)
+
         super().__init__(
             cache_dir=self.cache_build_dir,
             config_name=self.namespace,
             hash=sub_dir_hash,
-            data_files=self.meta_data_files,
+            data_files=data_files,
             **self.input_config_kwargs)
 
         self.info.builder_name = self.dataset_name
         self.name = camelcase_to_snakecase(self.dataset_name)
 
+        self.local_meta_csv_paths: dict = dict({})
+
     def _build_cache_dir(self, namespace=DEFAULT_DATASET_NAMESPACE):
         builder_data_dir = os.path.join(
             self._cache_dir_root,
@@ -147,6 +158,87 @@ class CsvDatasetBuilder(csv.Csv):
                     f"Failed to read file '{file}' with error {type(e)}: {e}")
                 raise
 
+    def download_and_prepare(self, download_mode, dl_manager,
+                             **download_kwargs):
+
+        target_cache_dir = dl_manager.download_config.cache_dir
+
+        split_name = dl_manager.download_config.split
+        if not split_name:
+            split_name = DatasetPathName.LOCK_FILE_NAME_ANY
+        version_name = dl_manager.download_config.version
+        if not version_name:
+            version_name = DatasetPathName.LOCK_FILE_NAME_ANY
+        subset_name = self.subset_name
+        if not subset_name:
+            subset_name = DatasetPathName.LOCK_FILE_NAME_ANY
+
+        # Prevent parallel disk operations
+        lock_file_names = []
+        lock_file_names.append(DatasetPathName.DATA_FILES_NAME)
+        lock_file_names.append(dl_manager.download_config.dataset_name)
+        lock_file_names.append(version_name)
+        lock_file_names.append(subset_name)
+        lock_file_names.append(split_name)
+
+        lock_file_name = DatasetPathName.LOCK_FILE_NAME_DELIMITER.join(
+            lock_file_names)
+
+        lock_path = os.path.join(
+            target_cache_dir.strip(DatasetPathName.DATA_FILES_NAME),
+            lock_file_name + '.lock')
+        with FileLock(lock_path):
+            data_exists = os.path.exists(target_cache_dir)
+            if data_exists and download_mode == DownloadMode.REUSE_DATASET_IF_EXISTS.value:
+                logger.warning(
+                    f'Reusing dataset {self.name} ({target_cache_dir})')
+            logger.info(f'Generating dataset {self.name} ({target_cache_dir})')
+
+            self._download_and_prepare(
+                dl_manager=dl_manager, download_mode=download_mode)
+
+    def _download_and_prepare(self, dl_manager, download_mode):
+        import shutil
+
+        target_cache_dir = dl_manager.download_config.cache_dir
+        if download_mode == DownloadMode.FORCE_REDOWNLOAD.value:
+            shutil.rmtree(target_cache_dir, ignore_errors=True)
+            os.makedirs(target_cache_dir, exist_ok=True)
+
+        self.local_meta_csv_paths = {
+            k: HubApi.fetch_csv_from_url(v, target_cache_dir)
+            for k, v in self.meta_data_files.items()
+        }
+
+        self.split_path_dict = dl_manager.download_and_extract(
+            self.zip_data_files)
+
+    def _convert_csv_to_dataset(self, split_name, csv_file_path):
+
+        df = pd.read_csv(
+            csv_file_path, iterator=False, delimiter=self.csv_delimiter)
+
+        transform_fields = []
+        for field_name in df.columns.tolist():
+            if field_name.endswith(':FILE'):
+                transform_fields.append(field_name)
+
+        base_extracted_dir = self.split_path_dict.get(split_name, '')
+        for field_name in transform_fields:
+            if base_extracted_dir:
+                df[field_name] = df[field_name].apply(
+                    lambda x: os.path.join(base_extracted_dir, x))
+
+        pa_data = pa.Table.from_pandas(df)
+        return Dataset(arrow_table=pa_data)
+
+    def as_dataset(self) -> DatasetDict:
+
+        return DatasetDict({
+            k: self._convert_csv_to_dataset(k, v)
+            for k, v in self.local_meta_csv_paths.items()
+        })
+
 
 class TaskSpecificDatasetBuilder(CsvDatasetBuilder):
 
@@ -181,7 +273,7 @@ class TaskSpecificDatasetBuilder(CsvDatasetBuilder):
             self._cache_dir.replace(os.sep, '_') + '.lock')
         with FileLock(lock_path):
             data_exists = os.path.exists(self._cache_dir)
-            if data_exists and download_mode == DownloadMode.REUSE_DATASET_IF_EXISTS:
+            if data_exists and download_mode == DownloadMode.REUSE_DATASET_IF_EXISTS:  # TODO: .value??
                 logger.warning(
                     f'Reusing dataset {self.name} ({self._cache_dir})')
                 return
@@ -233,6 +325,9 @@ class IterableDatasetBuilder(csv.Csv):
         self.info.builder_name = self.dataset_name
         self.name = camelcase_to_snakecase(self.dataset_name)
 
+        self.meta_csv_df = None
+        self.meta_cache_dir = dataset_context_config.data_meta_config.meta_cache_dir
+
     @staticmethod
     def get_builder_instance(
             dataset_context_config: DatasetContextConfig) -> csv.Csv:
@@ -357,17 +452,13 @@ class IterableDatasetBuilder(csv.Csv):
                 zip_file_name = os.path.splitext(zip_file)[0]
 
         if meta_file_url and not files:
-            headers, texts = hub_api.fetch_single_csv_script(meta_file_url)
-            meta_csv_mapping = IterableDatasetBuilder.trans_data_to_mapping(
-                headers, texts, self.csv_delimiter)
-            pa_table = pa.Table.from_pydict(meta_csv_mapping)
+            self._get_meta_csv_df(meta_file_url)
+            pa_table = pa.Table.from_pandas(self.meta_csv_df)
             yield 0, pa_table
 
         elif meta_file_url and files:
             # Get meta file
-            headers, texts = hub_api.fetch_single_csv_script(meta_file_url)
-            meta_csv_mapping = IterableDatasetBuilder.trans_data_to_mapping(
-                headers, texts, self.csv_delimiter)
+            self._get_meta_csv_df(meta_file_url)
 
             if is_zip:
                 oss_config_for_unzipped = hub_api.get_dataset_access_config_for_unzipped(
@@ -375,7 +466,7 @@ class IterableDatasetBuilder(csv.Csv):
                     zip_file_name)
                 dl_manager.download_config.oss_config = oss_config_for_unzipped
 
-            pa_table = pa.Table.from_pydict(meta_csv_mapping)
+            pa_table = pa.Table.from_pandas(self.meta_csv_df)
             yield 0, pa_table
 
         elif not meta_file_url and files:
@@ -385,6 +476,15 @@ class IterableDatasetBuilder(csv.Csv):
         else:
             raise f'Neither column meta nor data file found in {self.dataset_name}.json .'
 
+    def _get_meta_csv_df(self, meta_file_url: str) -> None:
+        if not self.meta_csv_df:
+            meta_csv_file_path = HubApi.fetch_csv_from_url(
+                meta_file_url, self.meta_cache_dir)
+            self.meta_csv_df = pd.read_csv(
+                meta_csv_file_path,
+                iterator=False,
+                delimiter=self.csv_delimiter)
+
     @staticmethod
     def trans_data_to_mapping(headers: str, texts: list, delimiter: str):
         res = {}

modelscope/msdatasets/download/download_config.py

@@ -15,6 +15,7 @@ class DataDownloadConfig(DownloadConfig):
         self.data_dir: Optional[str] = None
         self.oss_config: Optional[dict] = {}
         self.meta_args_map: Optional[dict] = {}
+        self.num_proc: int = 4
 
     def copy(self) -> 'DataDownloadConfig':
         return self

modelscope/msdatasets/meta/data_meta_manager.py

@@ -130,7 +130,8 @@ class DataMetaManager(object):
             target_subset_name, target_dataset_structure = get_target_dataset_structure(
                 dataset_json, subset_name, split)
             meta_map, file_map, args_map, type_map = get_dataset_files(
-                target_dataset_structure, dataset_name, namespace, version)
+                target_dataset_structure, dataset_name, namespace,
+                self.dataset_context_config, version)
 
             data_meta_config.meta_data_files = meta_map
             data_meta_config.zip_data_files = file_map

modelscope/msdatasets/ms_dataset.py

@@ -60,7 +60,8 @@ class MsDataset:
     _dataset_context_config: DatasetContextConfig = None
 
     def __init__(self,
-                 ds_instance: Union[Dataset, IterableDataset, ExternalDataset],
+                 ds_instance: Union[Dataset, IterableDataset, ExternalDataset,
+                                    NativeIterableDataset],
                  target: Optional[str] = None):
         self._hf_ds = ds_instance
         if target is not None and target not in self._hf_ds.features:
@@ -170,6 +171,7 @@ class MsDataset:
         cache_dir: Optional[str] = MS_DATASETS_CACHE,
         use_streaming: Optional[bool] = False,
         custom_cfg: Optional[Config] = Config(),
+        token: Optional[str] = None,
         **config_kwargs,
     ) -> Union[dict, 'MsDataset', NativeIterableDataset]:
         """Load a MsDataset from the ModelScope Hub, Hugging Face Hub, urls, or a local dataset.
@@ -197,12 +199,18 @@ class MsDataset:
                                                 NativeIterableDataset or a dict of NativeIterableDataset.
                 custom_cfg (str, Optional): Model configuration, this can be used for custom datasets.
                                            see https://modelscope.cn/docs/Configuration%E8%AF%A6%E8%A7%A3
+                token (str, Optional): SDK token of ModelScope.
                 **config_kwargs (additional keyword arguments): Keyword arguments to be passed
 
             Returns:
                 MsDataset (MsDataset): MsDataset object for a certain dataset.
             """
 
+        if token:
+            from modelscope.hub.api import HubApi
+            api = HubApi()
+            api.login(token)
+
         download_mode = DownloadMode(download_mode
                                      or DownloadMode.REUSE_DATASET_IF_EXISTS)
         hub = Hubs(hub or Hubs.modelscope)
@@ -647,7 +655,7 @@ class MsDataset:
 
             def type_converter(self, x):
                 if self.to_tensor:
-                    return torch.tensor(x)
+                    return torch.as_tensor(x)
                 else:
                     return x
 

modelscope/msdatasets/utils/dataset_utils.py

@@ -4,7 +4,11 @@ import os
 from collections import defaultdict
 from typing import Optional, Union
 
+import pandas as pd
+
 from modelscope.hub.api import HubApi
+from modelscope.msdatasets.context.dataset_context_config import \
+    DatasetContextConfig
 from modelscope.utils.constant import DEFAULT_DATASET_REVISION, MetaDataFields
 from modelscope.utils.logger import get_logger
 
@@ -169,6 +173,7 @@ def get_split_objects_map(file_map, objects):
 def get_dataset_files(subset_split_into: dict,
                       dataset_name: str,
                       namespace: str,
+                      context_config: DatasetContextConfig,
                       revision: Optional[str] = DEFAULT_DATASET_REVISION):
     """
     Return:
@@ -186,6 +191,7 @@ def get_dataset_files(subset_split_into: dict,
     args_map = defaultdict(dict)
     custom_type_map = defaultdict(dict)
     modelscope_api = HubApi()
+    meta_cache_dir = context_config.data_meta_config.meta_cache_dir
 
     for split, info in subset_split_into.items():
         custom_type_map[split] = info.get('custom', '')
@@ -200,16 +206,23 @@ def get_dataset_files(subset_split_into: dict,
     for split, args_dict in args_map.items():
         if args_dict and args_dict.get(MetaDataFields.ARGS_BIG_DATA):
             meta_csv_file_url = meta_map[split]
-            _, script_content = modelscope_api.fetch_single_csv_script(
-                meta_csv_file_url)
-            if not script_content:
-                raise 'Meta-csv file cannot be empty when meta-args `big_data` is true.'
-            for item in script_content:
-                if not item:
-                    continue
-                item = item.strip().split(',')[0]
-                if item:
-                    objects.append(item)
+
+            meta_csv_file_path = HubApi.fetch_csv_from_url(
+                meta_csv_file_url, meta_cache_dir)
+
+            csv_delimiter = context_config.config_kwargs.get('delimiter', ',')
+            csv_df = pd.read_csv(
+                meta_csv_file_path, iterator=False, delimiter=csv_delimiter)
+            target_col = csv_df.columns[csv_df.columns.str.contains(
+                ':FILE')].to_list()
+            if len(target_col) == 0:
+                logger.error(
+                    f'No column contains ":FILE" in {meta_csv_file_path}.')
+                target_col = csv_df.columns[0]
+            else:
+                target_col = target_col[0]
+            objects = csv_df[target_col].to_list()
+
             file_map[split] = objects
     # More general but low-efficiency.
     if not objects:

modelscope/msdatasets/utils/oss_utils.py

@@ -117,7 +117,8 @@ class OssUtilities:
                 if e.__dict__.get('status') == 403:
                     self._reload_sts()
                 if retry_count >= self.max_retries:
-                    raise
+                    logger.warning(f'Failed to download {oss_file_name}')
+                    raise e
 
         return local_path
 

modelscope/pipelines/audio/asr_inference_pipeline.py

@@ -127,7 +127,7 @@ class AutomaticSpeechRecognitionPipeline(Pipeline):
             minlenratio=self.cmd['minlenratio'],
             batch_size=self.cmd['batch_size'],
             beam_size=self.cmd['beam_size'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             ctc_weight=self.cmd['ctc_weight'],
             lm_weight=self.cmd['lm_weight'],
             penalty=self.cmd['penalty'],
@@ -387,8 +387,9 @@ class AutomaticSpeechRecognitionPipeline(Pipeline):
         ]
 
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/audio/lm_infer_pipeline.py

@@ -80,7 +80,7 @@ class LanguageModelPipeline(Pipeline):
             mode=self.cmd['mode'],
             batch_size=self.cmd['batch_size'],
             dtype=self.cmd['dtype'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             seed=self.cmd['seed'],
             num_workers=self.cmd['num_workers'],
             log_level=self.cmd['log_level'],
@@ -192,8 +192,9 @@ class LanguageModelPipeline(Pipeline):
         ]
 
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/audio/punctuation_processing_pipeline.py

@@ -54,7 +54,7 @@ class PunctuationProcessingPipeline(Pipeline):
             mode=self.cmd['mode'],
             batch_size=self.cmd['batch_size'],
             dtype=self.cmd['dtype'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             seed=self.cmd['seed'],
             num_workers=self.cmd['num_workers'],
             log_level=self.cmd['log_level'],
@@ -144,8 +144,9 @@ class PunctuationProcessingPipeline(Pipeline):
         ]
 
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/audio/speaker_diarization_pipeline.py

@@ -77,7 +77,7 @@ class SpeakerDiarizationPipeline(Pipeline):
             output_dir=self.cmd['output_dir'],
             batch_size=self.cmd['batch_size'],
             dtype=self.cmd['dtype'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             seed=self.cmd['seed'],
             num_workers=self.cmd['num_workers'],
             log_level=self.cmd['log_level'],
@@ -199,12 +199,13 @@ class SpeakerDiarizationPipeline(Pipeline):
 
         # rewrite the config with user args
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                if isinstance(cmd[user_args], dict) and isinstance(
-                        extra_args[user_args], dict):
-                    cmd[user_args].update(extra_args[user_args])
-                else:
-                    cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    if isinstance(cmd[user_args], dict) and isinstance(
+                            extra_args[user_args], dict):
+                        cmd[user_args].update(extra_args[user_args])
+                    else:
+                        cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/audio/speaker_verification_pipeline.py

@@ -57,7 +57,7 @@ class SpeakerVerificationPipeline(Pipeline):
             output_dir=self.cmd['output_dir'],
             batch_size=self.cmd['batch_size'],
             dtype=self.cmd['dtype'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             seed=self.cmd['seed'],
             num_workers=self.cmd['num_workers'],
             log_level=self.cmd['log_level'],
@@ -166,12 +166,13 @@ class SpeakerVerificationPipeline(Pipeline):
 
         # rewrite the config with user args
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                if isinstance(cmd[user_args], dict) and isinstance(
-                        extra_args[user_args], dict):
-                    cmd[user_args].update(extra_args[user_args])
-                else:
-                    cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    if isinstance(cmd[user_args], dict) and isinstance(
+                            extra_args[user_args], dict):
+                        cmd[user_args].update(extra_args[user_args])
+                    else:
+                        cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/audio/timestamp_pipeline.py

@@ -75,7 +75,7 @@ class TimestampPipeline(Pipeline):
             mode=self.cmd['mode'],
             batch_size=self.cmd['batch_size'],
             dtype=self.cmd['dtype'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             seed=self.cmd['seed'],
             num_workers=self.cmd['num_workers'],
             log_level=self.cmd['log_level'],
@@ -267,8 +267,9 @@ class TimestampPipeline(Pipeline):
         ]
 
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/audio/voice_activity_detection_pipeline.py

@@ -56,7 +56,7 @@ class VoiceActivityDetectionPipeline(Pipeline):
             mode=self.cmd['mode'],
             batch_size=self.cmd['batch_size'],
             dtype=self.cmd['dtype'],
-            ngpu=self.cmd['ngpu'],
+            ngpu=ngpu,
             seed=self.cmd['seed'],
             num_workers=self.cmd['num_workers'],
             log_level=self.cmd['log_level'],
@@ -212,8 +212,9 @@ class VoiceActivityDetectionPipeline(Pipeline):
         ]
 
         for user_args in user_args_dict:
-            if user_args in extra_args and extra_args[user_args] is not None:
-                cmd[user_args] = extra_args[user_args]
+            if user_args in extra_args:
+                if extra_args.get(user_args) is not None:
+                    cmd[user_args] = extra_args[user_args]
                 del extra_args[user_args]
 
         return cmd

modelscope/pipelines/builder.py

@@ -7,7 +7,8 @@ from modelscope.hub.snapshot_download import snapshot_download
 from modelscope.metainfo import DEFAULT_MODEL_FOR_PIPELINE, Pipelines
 from modelscope.models.base import Model
 from modelscope.utils.config import ConfigDict, check_config
-from modelscope.utils.constant import DEFAULT_MODEL_REVISION, Invoke
+from modelscope.utils.constant import (DEFAULT_MODEL_REVISION, Invoke,
+                                       ThirdParty)
 from modelscope.utils.hub import read_config
 from modelscope.utils.plugins import (register_modelhub_repo,
                                       register_plugins_repo)
@@ -18,7 +19,7 @@ from .util import is_official_hub_path
 PIPELINES = Registry('pipelines')
 
 
-def normalize_model_input(model, model_revision):
+def normalize_model_input(model, model_revision, third_party=None):
     """ normalize the input model, to ensure that a model str is a valid local path: in other words,
     for model represented by a model id, the model shall be downloaded locally
     """
@@ -26,19 +27,21 @@ def normalize_model_input(model, model_revision):
         # skip revision download if model is a local directory
         if not os.path.exists(model):
             # note that if there is already a local copy, snapshot_download will check and skip downloading
+            user_agent = {Invoke.KEY: Invoke.PIPELINE}
+            if third_party is not None:
+                user_agent[ThirdParty.KEY] = third_party
             model = snapshot_download(
-                model,
-                revision=model_revision,
-                user_agent={Invoke.KEY: Invoke.PIPELINE})
+                model, revision=model_revision, user_agent=user_agent)
     elif isinstance(model, list) and isinstance(model[0], str):
         for idx in range(len(model)):
             if is_official_hub_path(
                     model[idx],
                     model_revision) and not os.path.exists(model[idx]):
+                user_agent = {Invoke.KEY: Invoke.PIPELINE}
+                if third_party is not None:
+                    user_agent[ThirdParty.KEY] = third_party
                 model[idx] = snapshot_download(
-                    model[idx],
-                    revision=model_revision,
-                    user_agent={Invoke.KEY: Invoke.PIPELINE})
+                    model[idx], revision=model_revision, user_agent=user_agent)
     return model
 
 
@@ -97,7 +100,11 @@ def pipeline(task: str = None,
     if task is None and pipeline_name is None:
         raise ValueError('task or pipeline_name is required')
 
-    model = normalize_model_input(model, model_revision)
+    third_party = kwargs.get(ThirdParty.KEY)
+    if third_party is not None:
+        kwargs.pop(ThirdParty.KEY)
+    model = normalize_model_input(
+        model, model_revision, third_party=third_party)
     pipeline_props = {'type': pipeline_name}
     if pipeline_name is None:
         # get default pipeline for this task

modelscope/pipelines/cv/card_detection_pipeline.py

@@ -6,8 +6,8 @@ from modelscope.models.base.base_model import Model
 from modelscope.pipelines.base import Pipeline
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.utils.constant import Tasks
+from modelscope.utils.input_output_typing import Image
 from modelscope.utils.logger import get_logger
-from modelscope.utils.typing import Image
 
 logger = get_logger()
 

modelscope/pipelines/cv/face_detection_pipeline.py

@@ -16,8 +16,8 @@ from modelscope.pipelines.builder import PIPELINES
 from modelscope.preprocessors import LoadImage
 from modelscope.utils.config import Config
 from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.input_output_typing import Image
 from modelscope.utils.logger import get_logger
-from modelscope.utils.typing import Image
 
 logger = get_logger()
 

modelscope/pipelines/cv/ocr_recognition_pipeline.py

@@ -66,7 +66,7 @@ class OCRRecognitionPipeline(Pipeline):
         return outputs
 
     def forward(self, inputs):
-        outputs = self.ocr_recognizer(inputs)
+        outputs = self.ocr_recognizer(inputs['image'])
         return outputs
 
     def postprocess(self, inputs):

modelscope/pipelines/multi_modal/__init__.py

@@ -18,7 +18,7 @@ if TYPE_CHECKING:
     from .document_vl_embedding_pipeline import DocumentVLEmbeddingPipeline
     from .video_captioning_pipeline import VideoCaptioningPipeline
     from .video_question_answering_pipeline import VideoQuestionAnsweringPipeline
-    from .diffusers_wrapped import StableDiffusionWrapperPipeline, ChineseStableDiffusionPipeline
+    from .diffusers_wrapped import StableDiffusionPipeline, ChineseStableDiffusionPipeline
     from .soonet_video_temporal_grounding_pipeline import SOONetVideoTemporalGroundingPipeline
     from .text_to_video_synthesis_pipeline import TextToVideoSynthesisPipeline
     from .multimodal_dialogue_pipeline import MultimodalDialoguePipeline
@@ -42,7 +42,7 @@ else:
         'video_question_answering_pipeline':
         ['VideoQuestionAnsweringPipeline'],
         'diffusers_wrapped':
-        ['StableDiffusionWrapperPipeline', 'ChineseStableDiffusionPipeline'],
+        ['StableDiffusionPipeline', 'ChineseStableDiffusionPipeline'],
         'soonet_video_temporal_grounding_pipeline':
         ['SOONetVideoTemporalGroundingPipeline'],
         'text_to_video_synthesis_pipeline': ['TextToVideoSynthesisPipeline'],

modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py

@@ -4,12 +4,12 @@ from typing import TYPE_CHECKING
 from modelscope.utils.import_utils import LazyImportModule
 
 if TYPE_CHECKING:
-    from .stable_diffusion import StableDiffusionWrapperPipeline
+    from .stable_diffusion import StableDiffusionPipeline
     from .stable_diffusion import ChineseStableDiffusionPipeline
 else:
     _import_structure = {
         'stable_diffusion':
-        ['StableDiffusionWrapperPipeline', 'ChineseStableDiffusionPipeline']
+        ['StableDiffusionPipeline', 'ChineseStableDiffusionPipeline']
     }
 
     import sys

modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py

@@ -4,11 +4,11 @@ from typing import TYPE_CHECKING
 from modelscope.utils.import_utils import LazyImportModule
 
 if TYPE_CHECKING:
-    from .stable_diffusion_pipeline import StableDiffusionWrapperPipeline
+    from .stable_diffusion_pipeline import StableDiffusionPipeline
     from .chinese_stable_diffusion_pipeline import ChineseStableDiffusionPipeline
 else:
     _import_structure = {
-        'stable_diffusion_pipeline': ['StableDiffusionWrapperPipeline'],
+        'stable_diffusion_pipeline': ['StableDiffusionPipeline'],
         'chinese_stable_diffusion_pipeline':
         ['ChineseStableDiffusionPipeline']
     }

modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py

@@ -1,44 +1,49 @@
 # Copyright © Alibaba, Inc. and its affiliates.
-
-from typing import Any, Dict
+import os
+from typing import Any, Dict, Optional
 
 import cv2
 import numpy as np
 import torch
-from diffusers import StableDiffusionPipeline
+import torchvision.transforms as transforms
+from diffusers import \
+    StableDiffusionPipeline as DiffuserStableDiffusionPipeline
 from PIL import Image
 
 from modelscope.metainfo import Pipelines
+from modelscope.models import Model
 from modelscope.outputs import OutputKeys
+from modelscope.pipelines.base import Pipeline
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.pipelines.multi_modal.diffusers_wrapped.diffusers_pipeline import \
     DiffusersPipeline
 from modelscope.utils.constant import Tasks
 
 
-# Wrap around the diffusers stable diffusion pipeline implementation
-# for a unified ModelScope pipeline experience. Native stable diffusion
-# pipelines will be implemented in later releases.
 @PIPELINES.register_module(
     Tasks.text_to_image_synthesis,
     module_name=Pipelines.diffusers_stable_diffusion)
-class StableDiffusionWrapperPipeline(DiffusersPipeline):
+class StableDiffusionPipeline(DiffusersPipeline):
 
-    def __init__(self, model: str, device: str = 'gpu', **kwargs):
+    def __init__(self, model: str, lora_dir: str = None, **kwargs):
         """
         use `model` to create a stable diffusion pipeline
         Args:
-            model: model id on modelscope hub.
-            device: str = 'gpu'
+            model: model id on modelscope hub or local model dir.
         """
-        super().__init__(model, device, **kwargs)
 
-        torch_dtype = kwargs.get('torch_dtype', torch.float32)
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        # load pipeline
+        self.pipeline = DiffuserStableDiffusionPipeline.from_pretrained(
+            model, torch_dtype=torch.float16)
+        self.pipeline = self.pipeline.to(self.device)
+        # load lora moudle to unet
+        if lora_dir is not None:
+            assert os.path.exists(lora_dir), f"{lora_dir} isn't exist"
+            self.pipeline.unet.load_attn_procs(lora_dir)
 
-        # build upon the diffuser stable diffusion pipeline
-        self.pipeline = StableDiffusionPipeline.from_pretrained(
-            model, torch_dtype=torch_dtype)
-        self.pipeline.to(self.device)
+    def preprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        return inputs
 
     def forward(self, inputs: Dict[str, Any],
                 **forward_params) -> Dict[str, Any]:
@@ -46,24 +51,14 @@ class StableDiffusionWrapperPipeline(DiffusersPipeline):
             raise ValueError(
                 f'Expected the input to be a dictionary, but got {type(input)}'
             )
+
         if 'text' not in inputs:
             raise ValueError('input should contain "text", but not found')
 
-        return self.pipeline(
-            prompt=inputs.get('text'),
-            height=inputs.get('height'),
-            width=inputs.get('width'),
-            num_inference_steps=inputs.get('num_inference_steps', 50),
-            guidance_scale=inputs.get('guidance_scale', 7.5),
-            negative_prompt=inputs.get('negative_prompt'),
-            num_images_per_prompt=inputs.get('num_images_per_prompt', 1),
-            eta=inputs.get('eta', 0.0),
-            generator=inputs.get('generator'),
-            latents=inputs.get('latents'),
-            output_type=inputs.get('output_type', 'pil'),
-            return_dict=inputs.get('return_dict', True),
-            callback=inputs.get('callback'),
-            callback_steps=inputs.get('callback_steps', 1))
+        images = self.pipeline(
+            inputs['text'], num_inference_steps=30, guidance_scale=7.5)
+
+        return images
 
     def postprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
         images = []

modelscope/pipelines/nlp/distributed_gpt3_pipeline.py

@@ -1,6 +1,6 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
-from typing import Any, Dict
+from typing import Any, Dict, Generator, Optional
 
 import torch
 
@@ -9,12 +9,15 @@ from modelscope.models.nlp import DistributedGPT3
 from modelscope.pipelines.base import DistributedPipeline
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.preprocessors import TextGenerationJiebaPreprocessor
-from modelscope.utils.constant import Tasks
+from modelscope.utils.constant import Frameworks, Tasks
+from modelscope.utils.device import device_placement
+from modelscope.utils.streaming_output import PipelineStreamingOutputMixin
 
 
 @PIPELINES.register_module(
     Tasks.text_generation, module_name=Pipelines.gpt3_generation)
-class DistributedGPT3Pipeline(DistributedPipeline):
+class DistributedGPT3Pipeline(DistributedPipeline,
+                              PipelineStreamingOutputMixin):
     """This class is used to instantiate the gpt3 model.
     """
 
@@ -33,6 +36,8 @@ class DistributedGPT3Pipeline(DistributedPipeline):
             preprocessor = TextGenerationJiebaPreprocessor(model)
         super().__init__(model, preprocessor=preprocessor, **kwargs)
         assert hasattr(preprocessor, 'tokenizer')
+        self.model = PipelineStreamingOutputMixin()
+        self._model_prepare = True
 
     @classmethod
     def _instantiate_one(cls, rank, model_dir, **kwargs):
@@ -64,3 +69,36 @@ class DistributedGPT3Pipeline(DistributedPipeline):
 
     def _sanitize_parameters(self, **pipeline_parameters):
         return {}, pipeline_parameters, {}
+
+    def _stream_single(self, model_input: Dict[str, Any],
+                       forward_params: Dict[str, Any],
+                       postprocess_params: Dict[str, Any]) -> Generator:
+
+        with device_placement(self.framework, self.device_name):
+            if self._auto_collate:
+                model_input = self._collate_fn(model_input)
+            inputs = {'inputs': model_input, 'forward_params': forward_params}
+            self.model_pool.map(self.__class__._stream_one,
+                                [inputs] * self.world_size)
+
+        while True:
+            res = self.model_pool.map(self.__class__._next_one,
+                                      range(self.world_size))
+            if res[0] is None:
+                break
+            out = self.postprocess(res[0], **postprocess_params)
+            self._check_output(out)
+            yield out
+
+    @classmethod
+    def _stream_one(cls, inputs: Dict[str, Any]) -> None:
+        tokens = inputs['inputs']['input_ids'].cuda(
+            torch.cuda.current_device())
+        cls._stream = cls.model.stream(tokens, **inputs['forward_params'])
+
+    @classmethod
+    def _next_one(cls, idx: int) -> Optional[Dict[str, Any]]:
+        try:
+            return next(cls._stream)
+        except StopIteration:
+            return None

modelscope/pipelines/nlp/text_generation_pipeline.py

@@ -15,13 +15,14 @@ from modelscope.preprocessors import Preprocessor
 from modelscope.utils.chinese_utils import remove_space_between_chinese_chars
 from modelscope.utils.constant import ModelFile, Tasks
 from modelscope.utils.hub import Config, read_config
+from modelscope.utils.streaming_output import PipelineStreamingOutputMixin
 
 __all__ = ['TextGenerationPipeline', 'TextGenerationT5Pipeline']
 
 
 @PIPELINES.register_module(
     Tasks.text_generation, module_name=Pipelines.text_generation)
-class TextGenerationPipeline(Pipeline):
+class TextGenerationPipeline(Pipeline, PipelineStreamingOutputMixin):
 
     def __init__(self,
                  model: Union[Model, str],

modelscope/preprocessors/multi_modal.py

@@ -48,13 +48,14 @@ class DiffusionImageGenerationPreprocessor(Preprocessor):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.preprocessor_resolution = kwargs.pop('resolution', 512)
-        self.preprocessor_mean = kwargs.pop('mean', [0.5, 0.5, 0.5])
-        self.preprocessor_std = kwargs.pop('std', [0.5, 0.5, 0.5])
+        self.preprocessor_mean = kwargs.pop('mean', [0.5])
+        self.preprocessor_std = kwargs.pop('std', [0.5])
         self.preprocessor_image_keys = set(kwargs.pop('image_keys', []))
         self.transform_input = transforms.Compose([
             transforms.Resize(
                 self.preprocessor_resolution,
                 interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.RandomHorizontalFlip(),
             transforms.ToTensor(),
             transforms.Normalize(self.preprocessor_mean,
                                  self.preprocessor_std),

modelscope/preprocessors/nlp/token_classification_preprocessor.py

@@ -388,10 +388,14 @@ class TokenClassificationTransformersPreprocessor(
                 f'tokenizer {tokenizer_name}, please use a fast tokenizer instead, or '
                 f'try to implement a `{method}` method')
         label_mask, offset_mapping = getattr(self, method)(tokens)
-        padding = self.nlp_tokenizer.get_tokenizer_kwarg('padding')
-        max_length = self.nlp_tokenizer.get_tokenizer_kwarg('max_length')
-        special_token = 1 if self.nlp_tokenizer.get_tokenizer_kwarg(
-            'add_special_tokens') else 0
+        padding = kwargs.get('padding',
+                             self.nlp_tokenizer.get_tokenizer_kwarg('padding'))
+        max_length = kwargs.get(
+            'max_length', self.nlp_tokenizer.get_tokenizer_kwarg('max_length'))
+        special_token = 1 if kwargs.get(
+            'add_special_tokens',
+            self.nlp_tokenizer.get_tokenizer_kwarg(
+                'add_special_tokens')) else 0
         if len(label_mask) > max_length - 2 * special_token:
             label_mask = label_mask[:(max_length - 2 * special_token)]
         offset_mapping = offset_mapping[:sum(label_mask)]

modelscope/trainers/default_config.py

@@ -57,7 +57,7 @@ def update_cfg(cfg: Config) -> Config:
         key_chain_map[_HOOK_KEY_CHAIN_MAP[key]] = value
         hook.clear()
     cfg.train.hooks = list(filter(bool, cfg.train.hooks))
-    cfg.merge_from_dict(key_chain_map)
+    cfg.merge_from_dict(key_chain_map, force=False)
     return cfg
 
 

modelscope/trainers/hooks/checkpoint/checkpoint_hook.py

@@ -1,14 +1,16 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
 import random
-import time
+import shutil
 from typing import Optional
 
+import json
 import numpy as np
 import torch
 
 from modelscope.hub.check_model import check_model_is_id
-from modelscope.hub.push_to_hub import push_to_hub_async
+from modelscope.hub.push_to_hub import (UploadStrategy, push_to_hub_in_queue,
+                                        wait_for_done)
 from modelscope.metainfo import Hooks
 from modelscope.trainers.hooks.builder import HOOKS
 from modelscope.trainers.hooks.checkpoint.checkpoint_processor import \
@@ -45,7 +47,9 @@ class CheckpointHook(Hook):
         hub_repo_id (str): The hub repo id.
         hub_token (str): The token of the modelhub. You can also set the environment variable `MODELSCOPE_API_TOKEN`.
         private_hub (bool): Whether push to a private hub, default True.
-        hub_revision (str): Which branch to push the model to, default is `master`
+        hub_revision (str): Which branch to push the model to, default is `master`.
+        upload_strategy (str): The action adopted when the previous uploading is not done
+        and the next one is coming, can be `cancel` or `wait`.
         kwargs:
             by_epoch (bool): Same with `save_strategy`, but has a higher priority, legacy argument.
             output_sub_dir (str): The folder under the `save_dir` to save the output checkpoint for inference.
@@ -56,6 +60,8 @@ class CheckpointHook(Hook):
 
     EVAL_RESULT_FILE = 'eval_result.txt'
 
+    PUSH_TO_HUB_QUEUE_NAME = 'train.checkpoint'
+
     def __init__(self,
                  save_strategy: Optional[str] = CheckpointStrategy.by_epoch,
                  interval: Optional[int] = 0,
@@ -68,6 +74,7 @@ class CheckpointHook(Hook):
                  hub_token: Optional[str] = None,
                  private_hub: Optional[bool] = True,
                  hub_revision: Optional[str] = DEFAULT_REPOSITORY_REVISION,
+                 upload_strategy: Optional[str] = UploadStrategy.cancel,
                  **kwargs):
         self.interval = interval
         self.save_dir = save_dir
@@ -89,9 +96,9 @@ class CheckpointHook(Hook):
         self.hub_token = hub_token
         self.private_hub = private_hub
         self.hub_revision = hub_revision
+        self.upload_strategy = upload_strategy
         self.tag = -1
         self.is_model_id = None
-        self.push_to_hub_future = None
         self.max_checkpoint_num = None
         if max_checkpoint_num is not None:
             self.max_checkpoint_num = max(int(max_checkpoint_num), 1)
@@ -149,13 +156,15 @@ class CheckpointHook(Hook):
                         f'Saving checkpoint at {trainer.iter + 1} iter')
             self._save_checkpoint(trainer, prefix)
         if is_master() and self.push_to_hub:
-            if self.push_to_hub_future is not None and not self.push_to_hub_future.done(
-            ):
-                self.logger.error(
-                    f'Another uploading is running, '
-                    f'this uploading with message {prefix} will be canceled.')
-                return
-            self.push_to_hub_future = self._push_to_hub(trainer, prefix)
+            if self.upload_strategy == UploadStrategy.cancel:
+                output_dir = self.output_dir
+                delete_dir = False
+            else:
+                output_dir = self.output_dir + '_upload_' + prefix
+                shutil.copytree(
+                    self.output_dir, output_dir, dirs_exist_ok=True)
+                delete_dir = True
+            self._push_to_hub(trainer, prefix, output_dir, delete_dir)
 
     def after_train_epoch(self, trainer):
         if self.save_strategy != CheckpointStrategy.by_epoch:
@@ -172,32 +181,36 @@ class CheckpointHook(Hook):
             self._do_save(trainer, CheckpointStrategy.by_step)
 
     def after_run(self, trainer):
-        if self.push_to_hub_future is not None and not self.push_to_hub_future.done(
-        ):
-            self.logger.info('Train finished. Uploading models, waiting...')
-            while not self.push_to_hub_future.done():
-                time.sleep(1)
-            self.logger.info('Uploading models done.')
+        self.logger.info('Train finished. Uploading models, waiting...')
+        push_to_hub_in_queue(
+            self.PUSH_TO_HUB_QUEUE_NAME,
+            strategy=self.upload_strategy,
+            done=True)
+        wait_for_done(self.PUSH_TO_HUB_QUEUE_NAME)
+        self.logger.info('Uploading models done.')
 
-    def _push_to_hub(self, trainer, prefix):
+    def _push_to_hub(self, trainer, prefix, output_dir, delete_dir=False):
         if self.is_model_id is None:
             self.is_model_id = check_model_is_id(trainer.input_model_id,
                                                  self.hub_token)
         self.tag += 1
-        return push_to_hub_async(
-            self.hub_repo_id,
-            self.output_dir,
+        return push_to_hub_in_queue(
+            self.PUSH_TO_HUB_QUEUE_NAME,
+            strategy=self.upload_strategy,
+            repo_name=self.hub_repo_id,
+            output_dir=output_dir,
             token=self.hub_token,
             private=self.private_hub,
             commit_message=prefix,
             tag=f'v1.{self.tag}',
             revision=self.hub_revision,
-            source_repo=trainer.input_model_id if self.is_model_id else '')
+            source_repo=trainer.input_model_id if self.is_model_id else '',
+            delete_dir=delete_dir)
 
     def save_evaluate_results(self, trainer):
         with open(os.path.join(self.output_dir, self.EVAL_RESULT_FILE),
                   'w') as f:
-            f.write(str(trainer.metric_values))
+            f.write(json.dumps(trainer.metric_values))
 
     def _save_checkpoint(self, trainer, prefix):
         """Save checkpoint files and remove obsolete ones

modelscope/trainers/hooks/distributed/deepspeed_hook.py

@@ -1,206 +0,0 @@
-# Copyright (c) Alibaba, Inc. and its affiliates.
-import os
-import shutil
-
-import deepspeed
-import torch
-from deepspeed import DeepSpeedEngine
-from megatron_util import mpu, print_rank_0
-
-from modelscope.metainfo import Hooks
-from modelscope.trainers.hooks import LoadCheckpointHook
-from modelscope.trainers.hooks.builder import HOOKS
-from modelscope.trainers.hooks.checkpoint.checkpoint_hook import (
-    BestCkptSaverHook, CheckpointHook)
-from modelscope.trainers.hooks.hook import Hook
-from modelscope.trainers.hooks.priority import Priority
-from modelscope.utils.checkpoint import save_checkpoint
-from modelscope.utils.logger import get_logger
-from ..checkpoint.checkpoint_processor import CheckpointProcessor
-from ..lr_scheduler_hook import LrSchedulerProcessor
-from ..optimizer.base import OptimizerHook, OptimizerProcessor
-
-
-class DeepspeedProcessor(CheckpointProcessor, LrSchedulerProcessor,
-                         OptimizerProcessor):
-
-    _BIN_FILE_DIR = 'model'
-
-    def rank_name(self):
-        # TODO
-        try:
-            tp_world_size = mpu.get_tensor_model_parallel_world_size()
-            if tp_world_size == 1:
-                return ''
-            mp_rank = mpu.get_tensor_model_parallel_rank()
-            return '_mp_rank_{:02d}'.format(mp_rank)
-        except (ImportError, AssertionError):
-            return ''
-
-    def get_bin_file(self):
-        mp_rank = mpu.get_tensor_model_parallel_rank()
-        rank = '{:02d}'.format(mp_rank)
-        return f'mp_rank_{rank}_model_states.pt'
-
-    def save_checkpoints(self,
-                         trainer,
-                         checkpoint_path_prefix,
-                         output_dir,
-                         meta=None):
-        model = trainer.unwrap_module(trainer.model)
-        _train_state_file = checkpoint_path_prefix + self.rank_name(
-        ) + CheckpointProcessor.TRAINER_STATE_SUFFIX
-        # Save pth file without model state_dict
-        save_checkpoint(
-            model, _train_state_file, None, None, meta=meta, with_model=False)
-
-        save_dir = os.path.dirname(checkpoint_path_prefix)
-        prefix = os.path.basename(checkpoint_path_prefix)
-        trainer.model.save_checkpoint(save_dir, prefix)
-
-        bin_file = self.get_bin_file()
-        src_file = os.path.join(checkpoint_path_prefix, bin_file)
-        dest_file = os.path.join(output_dir, self._BIN_FILE_DIR, bin_file)
-        if os.path.isfile(dest_file):
-            os.unlink(dest_file)
-
-        try:
-            os.link(src_file, dest_file)
-        except OSError as e:
-            get_logger().error(
-                f'Link {src_file} to {dest_file} error: {e}, '
-                'changing to copy the bin file, this may case more space usage.'
-            )
-            shutil.copyfile(src_file, dest_file)
-
-    def remove_checkpoints(self, trainer, checkpoint_path_prefix):
-        _train_state_file = checkpoint_path_prefix + self.rank_name(
-        ) + CheckpointProcessor.TRAINER_STATE_SUFFIX
-        if os.path.isfile(_train_state_file):
-            os.remove(_train_state_file)
-
-        shutil.rmtree(checkpoint_path_prefix, ignore_errors=True)
-
-    def load_checkpoints(self, checkpoint_path_prefix, trainer, load_all_state,
-                         strict):
-        assert os.path.isdir(checkpoint_path_prefix)
-        path = os.path.dirname(checkpoint_path_prefix)
-        tag = os.path.basename(checkpoint_path_prefix)
-
-        meta = {}
-        _train_state_file = checkpoint_path_prefix + self.rank_name(
-        ) + CheckpointProcessor.TRAINER_STATE_SUFFIX
-        if os.path.isfile(_train_state_file):
-            meta = self.load_trainer_state(trainer, _train_state_file,
-                                           load_all_state)
-
-        if isinstance(trainer.model, DeepSpeedEngine):
-            # DeepSpeedEngine is initialized
-            trainer.model.load_checkpoint(
-                path,
-                tag,
-                load_module_strict=strict,
-                load_module_only=not load_all_state,
-            )
-        else:
-            # in eval or prediction
-            save_dir = checkpoint_path_prefix
-            bin_file = self.get_bin_file()
-            model_file = os.path.join(save_dir, bin_file)
-            checkpoint = torch.load(
-                model_file, map_location=lambda storage, loc: storage)
-            checkpoint = checkpoint['module']
-            model_dict = trainer.unwrap_module(trainer.model).state_dict()
-            for key in checkpoint:
-                if key not in model_dict.keys():
-                    print_rank_0('Skip key: ' + key)
-                else:
-                    print_rank_0('Loading key: ' + key)
-            trainer.unwrap_module(trainer.model).load_state_dict(
-                checkpoint, strict=strict)
-        return meta
-
-    def backward(self, trainer, loss_keys, cumulative_iters, grad_clip):
-        # assert cumulative_iters == 1, 'DeepSpeed only support cumulative_iters=1'
-        # The `trainer.model` here is actually a deepspeed engine object.
-        # backward step
-        for k in loss_keys:
-            loss = trainer.train_outputs[k]
-            trainer.model.backward(loss)
-
-        # update parameters
-        trainer.model.step()
-
-    def initialize_optimizer(self, trainer):
-        pass
-
-    def step(self, trainer):
-        pass
-
-
-@HOOKS.register_module(module_name=Hooks.DeepspeedHook)
-class DeepspeedHook(Hook):
-    PRIORITY = Priority.VERY_HIGH
-
-    def __init__(self,
-                 deepspeed_activation_checkpointing=True,
-                 save_zero_checkpoint=False,
-                 with_mpu=True):
-        self.save_zero_checkpoint = save_zero_checkpoint
-        self.deepspeed_activation_checkpointing = deepspeed_activation_checkpointing
-        # TODO without mpu
-        self.with_mpu = with_mpu
-        assert with_mpu, 'DeepspeedHook now is only for mpu models.'
-
-    def register_processor(self, trainer):
-        processor = DeepspeedProcessor()
-        optimizer_hook = trainer.get_hook(OptimizerHook)
-        if len(optimizer_hook) > 0 and not isinstance(
-                optimizer_hook[0].processor, DeepspeedProcessor):
-            optimizer_hook[0].set_processor(processor)
-        ckpt_hook = trainer.get_hook(CheckpointHook)
-        if len(ckpt_hook) > 0 and not isinstance(ckpt_hook[0].processor,
-                                                 DeepspeedProcessor):
-            ckpt_hook[0].set_processor(processor)
-        best_ckpt_hook = trainer.get_hook(BestCkptSaverHook)
-        if len(best_ckpt_hook) > 0 and not isinstance(
-                best_ckpt_hook[0].processor, DeepspeedProcessor):
-            best_ckpt_hook[0].set_processor(processor)
-        load_ckpt_hook = trainer.get_hook(LoadCheckpointHook)
-        if len(load_ckpt_hook) > 0 and not isinstance(
-                load_ckpt_hook[0].processor, DeepspeedProcessor):
-            load_ckpt_hook[0].set_processor(processor)
-
-    def before_val(self, trainer):
-        pass
-
-    def before_run(self, trainer):
-        if not hasattr(trainer, 'logger'):
-            self.logger = get_logger()
-        else:
-            self.logger = trainer.logger
-
-        # deepspeed init
-        args = trainer.cfg.train
-        args.deepspeed_config = os.path.join(trainer.model_dir,
-                                             args.deepspeed_config)
-
-        trainer.model, _, _, _ = deepspeed.initialize(
-            model=trainer.model,
-            optimizer=trainer.optimizer,
-            args=args,
-            lr_scheduler=trainer.lr_scheduler,
-            mpu=mpu,
-            dist_init_required=False)
-        trainer.model.save_zero_checkpoint = self.save_zero_checkpoint
-
-        if self.deepspeed_activation_checkpointing:
-            model = trainer.unwrap_module(trainer.model)
-            deepspeed.checkpointing.configure(
-                mpu,
-                deepspeed_config=args.deepspeed_config,
-                num_checkpoints=model.config.num_hidden_layers)
-
-            mpu.checkpoint = deepspeed.checkpointing.checkpoint
-            mpu.get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker
-            mpu.model_parallel_cuda_manual_seed = deepspeed.checkpointing.model_parallel_cuda_manual_seed

modelscope/trainers/hooks/lr_scheduler_hook.py

@@ -39,6 +39,20 @@ class LrSchedulerProcessor:
         else:
             trainer.lr_scheduler.step()
 
+    def get_current_lr(self, trainer):
+        import torch
+
+        if isinstance(trainer.optimizer, torch.optim.Optimizer):
+            lr = [group['lr'] for group in trainer.optimizer.param_groups]
+        elif isinstance(trainer.optimizer, dict):
+            lr = dict()
+            for name, optim in trainer.optimizer.items():
+                lr[name] = [group['lr'] for group in optim.param_groups]
+        else:
+            raise RuntimeError(
+                'lr is not applicable because optimizer does not exist.')
+        return lr
+
 
 class LrStrategy:
     by_epoch = 'by_epoch'
@@ -84,20 +98,6 @@ class LrSchedulerHook(Hook):
 
         self.processor.initialize_lr_scheduler(trainer)
 
-    def get_current_lr(self, trainer):
-        import torch
-
-        if isinstance(trainer.optimizer, torch.optim.Optimizer):
-            lr = [group['lr'] for group in trainer.optimizer.param_groups]
-        elif isinstance(trainer.optimizer, dict):
-            lr = dict()
-            for name, optim in trainer.optimizer.items():
-                lr[name] = [group['lr'] for group in optim.param_groups]
-        else:
-            raise RuntimeError(
-                'lr is not applicable because optimizer does not exist.')
-        return lr
-
     def after_train_iter(self, trainer):
         if self.lr_strategy == LrStrategy.by_step and trainer.iter >= getattr(
                 trainer, 'cumulative_iters', 1) - 1:
@@ -112,7 +112,7 @@ class LrSchedulerHook(Hook):
             self.processor.step(trainer)
 
     def _get_log_lr(self, trainer):
-        cur_lr = self.get_current_lr(trainer)
+        cur_lr = self.processor.get_current_lr(trainer)
         # only record lr of the first param group
         if isinstance(cur_lr, list):
             lr = cur_lr[0]

modelscope/trainers/multi_modal/lora_diffusion/__init__.py

@@ -0,0 +1 @@
+from .lora_diffusion_trainer import LoraDiffusionTrainer

modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py

@@ -0,0 +1,78 @@
+# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+from typing import Union
+
+import torch
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+
+from modelscope.metainfo import Trainers
+from modelscope.trainers.builder import TRAINERS
+from modelscope.trainers.hooks.checkpoint.checkpoint_hook import CheckpointHook
+from modelscope.trainers.hooks.checkpoint.checkpoint_processor import \
+    CheckpointProcessor
+from modelscope.trainers.optimizer.builder import build_optimizer
+from modelscope.trainers.trainer import EpochBasedTrainer
+from modelscope.utils.config import ConfigDict
+
+
+class LoraDiffusionCheckpointProcessor(CheckpointProcessor):
+
+    def save_checkpoints(self,
+                         trainer,
+                         checkpoint_path_prefix,
+                         output_dir,
+                         meta=None):
+        """Save the state dict for lora tune model.
+        """
+        trainer.model.unet = trainer.model.unet.to(torch.float32)
+        trainer.model.unet.save_attn_procs(output_dir)
+
+
+@TRAINERS.register_module(module_name=Trainers.lora_diffusion)
+class LoraDiffusionTrainer(EpochBasedTrainer):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # set lora save checkpoint processor
+        ckpt_hook = list(
+            filter(lambda hook: isinstance(hook, CheckpointHook),
+                   self.hooks))[0]
+        ckpt_hook.set_processor(LoraDiffusionCheckpointProcessor())
+        # Set correct lora layers
+        lora_attn_procs = {}
+        for name in self.model.unet.attn_processors.keys():
+            cross_attention_dim = None if name.endswith(
+                'attn1.processor'
+            ) else self.model.unet.config.cross_attention_dim
+            if name.startswith('mid_block'):
+                hidden_size = self.model.unet.config.block_out_channels[-1]
+            elif name.startswith('up_blocks'):
+                block_id = int(name[len('up_blocks.')])
+                hidden_size = list(
+                    reversed(
+                        self.model.unet.config.block_out_channels))[block_id]
+            elif name.startswith('down_blocks'):
+                block_id = int(name[len('down_blocks.')])
+                hidden_size = self.model.unet.config.block_out_channels[
+                    block_id]
+
+            lora_attn_procs[name] = LoRAAttnProcessor(
+                hidden_size=hidden_size,
+                cross_attention_dim=cross_attention_dim)
+
+        self.model.unet.set_attn_processor(lora_attn_procs)
+
+        self.lora_layers = AttnProcsLayers(self.model.unet.attn_processors)
+
+    def build_optimizer(self, cfg: ConfigDict, default_args: dict = None):
+        try:
+            return build_optimizer(
+                self.lora_layers.parameters(),
+                cfg=cfg,
+                default_args=default_args)
+        except KeyError as e:
+            self.logger.error(
+                f'Build optimizer error, the optimizer {cfg} is a torch native component, '
+                f'please check if your torch with version: {torch.__version__} matches the config.'
+            )
+            raise e

modelscope/trainers/multi_modal/stable_diffusion/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+from .stable_diffusion_trainer import StableDiffusionTrainer

modelscope/trainers/multi_modal/stable_diffusion/stable_diffusion_trainer.py

@@ -0,0 +1,30 @@
+# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+from typing import Union
+
+import torch
+from torch import nn
+
+from modelscope.metainfo import Trainers
+from modelscope.models.base import Model, TorchModel
+from modelscope.trainers.builder import TRAINERS
+from modelscope.trainers.optimizer.builder import build_optimizer
+from modelscope.trainers.trainer import EpochBasedTrainer
+from modelscope.utils.config import ConfigDict
+
+
+@TRAINERS.register_module(module_name=Trainers.stable_diffusion)
+class StableDiffusionTrainer(EpochBasedTrainer):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def build_optimizer(self, cfg: ConfigDict, default_args: dict = None):
+        try:
+            return build_optimizer(
+                self.model.unet, cfg=cfg, default_args=default_args)
+        except KeyError as e:
+            self.logger.error(
+                f'Build optimizer error, the optimizer {cfg} is a torch native component, '
+                f'please check if your torch with version: {torch.__version__} matches the config.'
+            )
+            raise e

modelscope/trainers/trainer.py

@@ -1001,7 +1001,7 @@ class EpochBasedTrainer(BaseTrainer):
         """
         optimizer, lr_scheduler = self.optimizers
         if optimizer is None:
-            optimizer_cfg = self.cfg.train.get('optimizer', None)
+            optimizer_cfg = deepcopy(self.cfg.train.get('optimizer', None))
         else:
             optimizer_cfg = None
 
@@ -1011,7 +1011,8 @@ class EpochBasedTrainer(BaseTrainer):
             optimizer = self.build_optimizer(cfg=optimizer_cfg)
 
         if lr_scheduler is None:
-            lr_scheduler_cfg = self.cfg.train.get('lr_scheduler', None)
+            lr_scheduler_cfg = deepcopy(
+                self.cfg.train.get('lr_scheduler', None))
         else:
             lr_scheduler_cfg = None
 

modelscope/trainers/training_args.py

@@ -9,6 +9,7 @@ import json
 
 from modelscope.trainers.cli_argument_parser import CliArgumentParser
 from modelscope.utils.config import Config
+from modelscope.utils.constant import DEFAULT_DATASET_NAMESPACE
 
 
 def set_flatten_value(values: Union[str, List[str]]):
@@ -62,13 +63,13 @@ class DatasetArgs:
         })
 
     train_dataset_namespace: str = field(
-        default=None,
+        default=DEFAULT_DATASET_NAMESPACE,
         metadata={
             'help': 'The dataset namespace used for training',
         })
 
     val_dataset_namespace: str = field(
-        default=None,
+        default=DEFAULT_DATASET_NAMESPACE,
         metadata={
             'help': 'The dataset namespace used for evaluating',
         })
@@ -116,6 +117,11 @@ class ModelArgs:
             'help': 'A model id or model dir',
         })
 
+    model_revision: str = field(
+        default=None, metadata={
+            'help': 'the revision of model',
+        })
+
     model_type: str = field(
         default=None,
         metadata={
@@ -450,7 +456,7 @@ class TrainingArgs(DatasetArgs, TrainArgs, ModelArgs):
             _unknown[unknown[i].replace('-', '')] = parse_value(unknown[i + 1])
         args_dict = vars(args)
         self.manual_args += parser.manual_args
-
+        self._unknown_args.update(_unknown)
         for key, value in deepcopy(args_dict).items():
             if key is not None and hasattr(self, key):
                 setattr(self, key, value)
@@ -506,7 +512,7 @@ def build_dataset_from_file(filename):
                 "text2": "sequence2",
                 "label": "label",
             }
-            "split": 0.8,
+            "usage": 0.8,
         }
     ]
     """
@@ -540,16 +546,16 @@ def build_dataset_from_file(filename):
                 lambda x: x,
                 remove_columns=remove_columns,
                 features=new_features).rename_columns(ds['column_mapping'])
-            split = ds['split']
-            if isinstance(split, str):
-                assert split in ('train', 'val')
-                if split == 'train':
+            usage = ds['usage']
+            if isinstance(usage, str):
+                assert usage in ('train', 'val')
+                if usage == 'train':
                     train_set.append(dataset)
                 else:
                     eval_set.append(dataset)
             else:
-                assert isinstance(split, float) and 0 < split < 1
-                ds_dict = dataset.train_test_split(train_size=split)
+                assert isinstance(usage, float) and 0 < usage < 1
+                ds_dict = dataset.train_test_split(train_size=usage)
                 train_set.append(ds_dict['train'])
                 eval_set.append(ds_dict['test'])
 

modelscope/tuners/control_sd_lora.py

@@ -12,8 +12,8 @@ import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.cross_attention import CrossAttention, LoRALinearLayer
 from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.resnet import (Downsample2D, Mish, Upsample2D,
-                                     downsample_2d, partial, upsample_2d)
+from diffusers.models.resnet import (Downsample2D, Upsample2D, downsample_2d,
+                                     partial, upsample_2d)
 from diffusers.models.unet_2d_blocks import \
     get_down_block as get_down_block_default
 from diffusers.utils.outputs import BaseOutput
@@ -477,8 +477,10 @@ class ControlLoRACrossAttnProcessor(LoRACrossAttnProcessor):
         assert self.control_states is not None
 
         batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(attention_mask,
-                                                     sequence_length)
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask=attention_mask,
+            target_length=sequence_length,
+            batch_size=batch_size)
         query = attn.to_q(hidden_states)
         for pre_lora in self.pre_loras:
             lora_in = query if pre_lora.post_add else hidden_states
@@ -627,8 +629,10 @@ class ControlLoRACrossAttnProcessorV2(LoRACrossAttnProcessor):
         assert self.control_states is not None
 
         batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(attention_mask,
-                                                     sequence_length)
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask=attention_mask,
+            target_length=sequence_length,
+            batch_size=batch_size)
         for pre_lora in self.pre_loras:
             if isinstance(pre_lora, ControlLoRACrossAttnProcessorV2):
                 hidden_states = hidden_states + pre_lora.process_control_states(
@@ -783,7 +787,7 @@ class ConvBlock2D(nn.Module):
         if non_linearity == 'swish':
             self.nonlinearity = lambda x: F.silu(x)
         elif non_linearity == 'mish':
-            self.nonlinearity = Mish()
+            self.nonlinearity = nn.Mish()
         elif non_linearity == 'silu':
             self.nonlinearity = nn.SiLU()
 

modelscope/tuners/sd_lora.py

@@ -90,9 +90,10 @@ class LoRACrossAttnProcessor(nn.Module):
                  attention_mask=None,
                  scale=1.0):
         batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(attention_mask,
-                                                     sequence_length,
-                                                     batch_size)
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask=attention_mask,
+            target_length=sequence_length,
+            batch_size=batch_size)
 
         query = attn.to_q(hidden_states)
         query = query + scale * self.to_q_lora(

modelscope/utils/plugins.py

@@ -171,7 +171,21 @@ def import_plugins(plugins: List[str] = None) -> List[str]:
 
     for module_name in plugins:
         try:
-            import_module_and_submodules(module_name)
+            # TODO: include and exclude should be configurable, hard code now
+            import_module_and_submodules(
+                module_name,
+                include={
+                    'easycv.toolkit.modelscope',
+                    'easycv.hooks',
+                    'easycv.models',
+                    'easycv.core',
+                    'easycv.toolkit',
+                    'easycv.predictors',
+                },
+                exclude={
+                    'easycv.toolkit.*',
+                    'easycv.*',
+                })
             logger.info('Plugin %s available', module_name)
             imported_plugins.append(module_name)
         except ModuleNotFoundError as e:
@@ -238,9 +252,10 @@ def import_module_and_submodules(package_name: str,
             path_string = '' if not path else path[0]
 
             # walk_packages only finds immediate children, so need to recurse.
-            for module_finder, name, _ in pkgutil.walk_packages(path):
+            for module_finder, name, _ in pkgutil.iter_modules(path):
                 # Sometimes when you import third-party libraries that are on your path,
                 # `pkgutil.walk_packages` returns those too, so we need to skip them.
+                # `pkgutil.iter_modules` avoid import those package
                 if path_string and module_finder.path != path_string:  # type: ignore[union-attr]
                     continue
                 if name.startswith('_'):
@@ -250,7 +265,8 @@ def import_module_and_submodules(package_name: str,
                     # skip tests
                     continue
                 subpackage = f'{package_name}.{name}'
-                import_module_and_submodules(subpackage, exclude=exclude)
+                import_module_and_submodules(
+                    subpackage, include=include, exclude=exclude)
         except SystemExit as e:
             # this case is specific for easy_cv's tools/predict.py exit
             logger.warning(

modelscope/utils/pre_compile.py

@@ -0,0 +1,27 @@
+import os
+
+import torch
+
+from modelscope.utils.megatron_utils import init_megatron_util
+
+
+def pre_compile_megatron_util():
+    dummy_megatron_cfg = {
+        'tensor_model_parallel_size': 1,
+        'world_size': 1,
+        'distributed_backend': 'nccl',
+        'seed': 42,
+    }
+    os.environ['MASTER_PORT'] = '39501'
+    init_megatron_util(dummy_megatron_cfg)
+
+
+def pre_compile_all():
+    if torch.cuda.is_available():  # extension require cuda.
+        pre_compile_megatron_util()
+
+    # extension for all platform.
+
+
+if __name__ == '__main__':
+    pre_compile_all()

modelscope/utils/streaming_output.py

@@ -0,0 +1,145 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict, Generator, List, Union
+
+import torch
+
+from modelscope.pipelines.base import Input
+from modelscope.utils.constant import Frameworks
+from modelscope.utils.device import device_placement
+
+
+class StreamingOutputMixin:
+
+    def stream(self, *args, **kwargs) -> Generator:
+        """
+        Support the input of Model and Pipeline.
+        The output is a `Generator` type,
+        which conforms to the output standard of modelscope.
+        """
+        raise NotImplementedError
+
+
+class PipelineStreamingOutputMixin(StreamingOutputMixin):
+
+    def stream(self, input: Union[Input, List[Input]], *args,
+               **kwargs) -> Generator:
+        """
+        Similar to the `Pipeline.__call__` method.
+        it supports the input that the pipeline can accept,
+        and also supports batch input.
+
+        self.model must be a subclass of StreamingOutputMixin
+        and implement the stream method.
+        """
+        assert isinstance(self.model, StreamingOutputMixin
+                          ), 'pipeline.model must be StreamingOutputMixin!'
+        if (self.model or (self.has_multiple_models and self.models[0])):
+            if not self._model_prepare:
+                self.prepare_model()
+
+        batch_size = kwargs.pop('batch_size', None)
+        preprocess_params, forward_params, postprocess_params = self._sanitize_parameters(
+            **kwargs)
+
+        if isinstance(input, list):
+            model_input_list = [
+                self._preprocess_with_check(i, preprocess_params)
+                for i in input
+            ]
+
+            if batch_size is None:
+                output = []
+                for ele in model_input_list:
+                    output.append(
+                        self._stream_single(ele, forward_params,
+                                            postprocess_params))
+            else:
+                output = self._stream_batch(model_input_list, batch_size,
+                                            forward_params, postprocess_params)
+
+        else:
+            model_input = self._preprocess_with_check(input, preprocess_params)
+            output = self._stream_single(model_input, forward_params,
+                                         postprocess_params)
+        return output
+
+    def _preprocess_with_check(
+            self, input: Input,
+            preprocess_params: Dict[str, Any]) -> Dict[str, Any]:
+        self._check_input(input)
+        return self.preprocess(input, **preprocess_params)
+
+    def _stream_single(self, model_input: Dict[str, Any],
+                       forward_params: Dict[str, Any],
+                       postprocess_params: Dict[str, Any]) -> Generator:
+
+        with device_placement(self.framework, self.device_name):
+            if self.framework == Frameworks.torch:
+                with torch.no_grad():
+                    if self._auto_collate:
+                        model_input = self._collate_fn(model_input)
+                    stream = self.model.stream(model_input, **forward_params)
+            else:
+                stream = self.model.stream(model_input, **forward_params)
+
+            for out in stream:
+                out = self.postprocess(out, **postprocess_params)
+                self._check_output(out)
+                yield out
+
+    def _stream_batch(self, model_input_list: List[Dict[str, Any]],
+                      batch_size: int, forward_params: Dict[str, Any],
+                      postprocess_params: Dict[str, Any]) -> Generator:
+
+        stream_list = []
+        real_batch_sizes = []
+        with device_placement(self.framework, self.device_name):
+            for i in range(0, len(model_input_list), batch_size):
+                end = min(i + batch_size, len(model_input_list))
+                real_batch_size = end - i
+                real_batch_sizes.append(real_batch_size)
+
+                batched_out = self._batch(model_input_list[i:end])
+                if self.framework == Frameworks.torch:
+                    with torch.no_grad():
+                        if self._auto_collate:
+                            batched_out = self._collate_fn(batched_out)
+                        stream_list.append(
+                            self.model.stream(batched_out, **forward_params))
+                else:
+                    stream_list.append(
+                        self.model.stream(batched_out, **forward_params))
+
+            output_list = [None] * len(model_input_list)
+            stop_streams = 0
+            while stop_streams < len(stream_list):
+                stop_streams = 0
+                for i, (stream, real_batch_size) in enumerate(
+                        zip(stream_list, real_batch_sizes)):
+                    try:
+                        batched_out = next(stream)
+                        for batch_idx in range(real_batch_size):
+                            out = {}
+                            for k, element in batched_out.items():
+                                if element is not None:
+                                    if isinstance(element, (tuple, list)):
+                                        if isinstance(element[0],
+                                                      torch.Tensor):
+                                            out[k] = type(element)(
+                                                e[batch_idx:batch_idx + 1]
+                                                for e in element)
+                                        else:
+                                            # Compatible with traditional pipelines
+                                            out[k] = element[batch_idx]
+                                    else:
+                                        out[k] = element[batch_idx:batch_idx
+                                                         + 1]
+                            out = self.postprocess(out, **postprocess_params)
+                            self._check_output(out)
+                            output_index = i * batch_size + batch_idx
+                            output_list[output_index] = out
+                    except StopIteration:
+                        stop_streams += 1
+                yield output_list
+
+        return output_list

modelscope/version.py

@@ -1,5 +1,5 @@
 # Make sure to modify __release_datetime__ to release time when making official release.
-__version__ = '1.6.0'
+__version__ = '1.6.2'
 # default release datetime for branches under active development is set
 # to be a time far-far-away-into-the-future
 __release_datetime__ = '2099-10-13 08:56:12'