diff --git a/.dev_scripts/build_base_image.sh b/.dev_scripts/build_base_image.sh new file mode 100644 index 00000000..e60f6b3c --- /dev/null +++ b/.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 diff --git a/.dev_scripts/build_image.sh b/.dev_scripts/build_image.sh index 6bc8b5e4..a9cbef58 100644 --- a/.dev_scripts/build_image.sh +++ b/.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" diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index 8d1c9a0c..931ca397 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.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 diff --git a/.gitignore b/.gitignore index 00771ab4..6cc2df63 100644 --- a/.gitignore +++ b/.gitignore @@ -129,6 +129,5 @@ result.mp4 *.pth *.pt - # ast template ast_index_file.py diff --git a/data/test b/data/test index d0cd4f21..acc59489 160000 --- a/data/test +++ b/data/test @@ -1 +1 @@ -Subproject commit d0cd4f218f200267a7b24ae38df53e47cb15bf2d +Subproject commit acc59489d3954fc09cd99d8d4aed818a8d39b283 diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu index 27c0f1f3..2579a6de 100644 --- a/docker/Dockerfile.ubuntu +++ b/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 diff --git a/docker/Dockerfile.ubuntu_base b/docker/Dockerfile.ubuntu_base new file mode 100644 index 00000000..4d92ceda --- /dev/null +++ b/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 diff --git a/docker/rcfiles/sources.list.aliyun b/docker/rcfiles/sources.list.aliyun index 120bb1f1..1ebf4ae5 100644 --- a/docker/rcfiles/sources.list.aliyun +++ b/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 diff --git a/docker/scripts/install_apex.sh b/docker/scripts/install_apex.sh index 47f34da7..40d9f268 100644 --- a/docker/scripts/install_apex.sh +++ b/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 diff --git a/docs/source/api/modelscope.pipelines.multi_modal.rst b/docs/source/api/modelscope.pipelines.multi_modal.rst index b62878e3..ffb65fad 100644 --- a/docs/source/api/modelscope.pipelines.multi_modal.rst +++ b/docs/source/api/modelscope.pipelines.multi_modal.rst @@ -18,7 +18,7 @@ modelscope.pipelines.multi_modal ImageCaptioningPipeline MGeoRankingPipeline MultiModalEmbeddingPipeline - StableDiffusionWrapperPipeline + StableDiffusionPipeline TextToImageSynthesisPipeline VideoCaptioningPipeline VideoMultiModalEmbeddingPipeline diff --git a/examples/pytorch/llama/finetune_llama.py b/examples/pytorch/llama/finetune_llama.py new file mode 100644 index 00000000..88975e66 --- /dev/null +++ b/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 = '' +DEFAULT_BOS_TOKEN = '' +DEFAULT_UNK_TOKEN = '' +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() diff --git a/examples/pytorch/llama/run_train_llama.sh b/examples/pytorch/llama/run_train_llama.sh new file mode 100644 index 00000000..7c860d57 --- /dev/null +++ b/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 diff --git a/examples/pytorch/named_entity_recognition/finetune_named_entity_recognition.py b/examples/pytorch/named_entity_recognition/finetune_named_entity_recognition.py new file mode 100644 index 00000000..67a091f7 --- /dev/null +++ b/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() diff --git a/examples/pytorch/named_entity_recognition/run_train.sh b/examples/pytorch/named_entity_recognition/run_train.sh new file mode 100644 index 00000000..91b2c5bf --- /dev/null +++ b/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 \ diff --git a/examples/pytorch/stable_diffusion/finetune_stable_diffusion.py b/examples/pytorch/stable_diffusion/finetune_stable_diffusion.py index 28ba853c..ac16ed2c 100644 --- a/examples/pytorch/stable_diffusion/finetune_stable_diffusion.py +++ b/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() diff --git a/examples/pytorch/stable_diffusion/run_train.sh b/examples/pytorch/stable_diffusion/run_train.sh index 0e551942..8e45ae88 100644 --- a/examples/pytorch/stable_diffusion/run_train.sh +++ b/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 diff --git a/examples/pytorch/text_classification/finetune_text_classification.py b/examples/pytorch/text_classification/finetune_text_classification.py index dfcb7b4d..0a53fb7f 100644 --- a/examples/pytorch/text_classification/finetune_text_classification.py +++ b/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, diff --git a/examples/pytorch/text_classification/run_train.sh b/examples/pytorch/text_classification/run_train.sh index e91a9996..e148655c 100644 --- a/examples/pytorch/text_classification/run_train.sh +++ b/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' \ diff --git a/examples/pytorch/text_generation/finetune_text_generation.py b/examples/pytorch/text_generation/finetune_text_generation.py index a89970e8..588d83f3 100644 --- a/examples/pytorch/text_generation/finetune_text_generation.py +++ b/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) diff --git a/examples/pytorch/text_generation/run_train_gpt3.sh b/examples/pytorch/text_generation/run_train_gpt3.sh index fd37b42c..78521951 100644 --- a/examples/pytorch/text_generation/run_train_gpt3.sh +++ b/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' \ diff --git a/examples/pytorch/text_generation/run_train_mt5.sh b/examples/pytorch/text_generation/run_train_mt5.sh index 6d032d6e..b2d0bbf1 100644 --- a/examples/pytorch/text_generation/run_train_mt5.sh +++ b/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 \ diff --git a/examples/pytorch/text_generation/run_train_palm.sh b/examples/pytorch/text_generation/run_train_palm.sh index 68b9e89d..06153812 100644 --- a/examples/pytorch/text_generation/run_train_palm.sh +++ b/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 \ diff --git a/modelscope/exporters/cv/object_detection_damoyolo_exporter.py b/modelscope/exporters/cv/object_detection_damoyolo_exporter.py index 673811ad..6f271d39 100644 --- a/modelscope/exporters/cv/object_detection_damoyolo_exporter.py +++ b/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): diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index e3436aea..2f14976d 100644 --- a/modelscope/hub/api.py +++ b/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, diff --git a/modelscope/hub/push_to_hub.py b/modelscope/hub/push_to_hub.py index d117cc7f..2b2b4091 100644 --- a/modelscope/hub/push_to_hub.py +++ b/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) diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index c4057314..f2529be2 100644 --- a/modelscope/metainfo.py +++ b/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): diff --git a/modelscope/models/cv/image_human_parsing/parsing_utils.py b/modelscope/models/cv/image_human_parsing/parsing_utils.py index a1c20072..4d1652b4 100644 --- a/modelscope/models/cv/image_human_parsing/parsing_utils.py +++ b/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: diff --git a/modelscope/models/cv/ocr_detection/model.py b/modelscope/models/cv/ocr_detection/model.py index 712973ce..5e1728bf 100644 --- a/modelscope/models/cv/ocr_detection/model.py +++ b/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}' diff --git a/modelscope/models/cv/ocr_detection/modules/dbnet.py b/modelscope/models/cv/ocr_detection/modules/dbnet.py index 82b0e512..2745dbf4 100644 --- a/modelscope/models/cv/ocr_detection/modules/dbnet.py +++ b/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): diff --git a/modelscope/models/cv/ocr_detection/modules/layers.py b/modelscope/models/cv/ocr_detection/modules/layers.py new file mode 100644 index 00000000..515c0bc6 --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/ocr_detection/modules/mix_ops.py b/modelscope/models/cv/ocr_detection/modules/mix_ops.py new file mode 100644 index 00000000..5c6dd4a5 --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/ocr_detection/modules/proxyless.py b/modelscope/models/cv/ocr_detection/modules/proxyless.py new file mode 100644 index 00000000..68458bc3 --- /dev/null +++ b/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) diff --git a/modelscope/models/cv/ocr_recognition/model.py b/modelscope/models/cv/ocr_recognition/model.py index 2406b6dc..3510de45 100644 --- a/modelscope/models/cv/ocr_recognition/model.py +++ b/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) diff --git a/modelscope/models/cv/ocr_recognition/modules/CRNN/__init__.py b/modelscope/models/cv/ocr_recognition/modules/CRNN/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/ocr_recognition/modules/crnn.py b/modelscope/models/cv/ocr_recognition/modules/CRNN/main_model.py similarity index 92% rename from modelscope/models/cv/ocr_recognition/modules/crnn.py rename to modelscope/models/cv/ocr_recognition/modules/CRNN/main_model.py index 3de8304d..f2ce8c21 100644 --- a/modelscope/models/cv/ocr_recognition/modules/crnn.py +++ b/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) diff --git a/modelscope/models/cv/ocr_recognition/modules/ConvNextViT/__init__.py b/modelscope/models/cv/ocr_recognition/modules/ConvNextViT/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/ocr_recognition/modules/convnext.py b/modelscope/models/cv/ocr_recognition/modules/ConvNextViT/convnext.py similarity index 100% rename from modelscope/models/cv/ocr_recognition/modules/convnext.py rename to modelscope/models/cv/ocr_recognition/modules/ConvNextViT/convnext.py diff --git a/modelscope/models/cv/ocr_recognition/modules/convnextvit.py b/modelscope/models/cv/ocr_recognition/modules/ConvNextViT/main_model.py similarity index 63% rename from modelscope/models/cv/ocr_recognition/modules/convnextvit.py rename to modelscope/models/cv/ocr_recognition/modules/ConvNextViT/main_model.py index 4e7900b1..365d4005 100644 --- a/modelscope/models/cv/ocr_recognition/modules/convnextvit.py +++ b/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 diff --git a/modelscope/models/cv/ocr_recognition/modules/timm_tinyc.py b/modelscope/models/cv/ocr_recognition/modules/ConvNextViT/timm_tinyc.py similarity index 100% rename from modelscope/models/cv/ocr_recognition/modules/timm_tinyc.py rename to modelscope/models/cv/ocr_recognition/modules/ConvNextViT/timm_tinyc.py diff --git a/modelscope/models/cv/ocr_recognition/modules/vitstr.py b/modelscope/models/cv/ocr_recognition/modules/ConvNextViT/vitstr.py similarity index 100% rename from modelscope/models/cv/ocr_recognition/modules/vitstr.py rename to modelscope/models/cv/ocr_recognition/modules/ConvNextViT/vitstr.py diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/__init__.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py new file mode 100644 index 00000000..08584b5c --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/__init__.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/__init__.py new file mode 100644 index 00000000..c8115b7c --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/__init__.py @@ -0,0 +1 @@ +from .proxyless import plnas_linear_mix_se diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/layers.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/layers.py new file mode 100644 index 00000000..23802f6d --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/mix_ops.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/mix_ops.py new file mode 100644 index 00000000..effed1a6 --- /dev/null +++ b/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] diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py new file mode 100644 index 00000000..4f525639 --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/ocr_recognition/preprocessor.py b/modelscope/models/cv/ocr_recognition/preprocessor.py index d327cd3c..b111637f 100644 --- a/modelscope/models/cv/ocr_recognition/preprocessor.py +++ b/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} diff --git a/modelscope/models/cv/table_recognition/lineless_table_process.py b/modelscope/models/cv/table_recognition/lineless_table_process.py index 0d7fcfb5..a04d531d 100644 --- a/modelscope/models/cv/table_recognition/lineless_table_process.py +++ b/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 diff --git a/modelscope/models/multi_modal/stable_diffusion/__init__.py b/modelscope/models/multi_modal/stable_diffusion/__init__.py new file mode 100644 index 00000000..9ffff12a --- /dev/null +++ b/modelscope/models/multi_modal/stable_diffusion/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from .stable_diffusion import StableDiffusion diff --git a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py new file mode 100644 index 00000000..8ec2149d --- /dev/null +++ b/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) diff --git a/modelscope/models/nlp/gpt3/backbone.py b/modelscope/models/nlp/gpt3/backbone.py index 2f8e4699..89e19480 100644 --- a/modelscope/models/nlp/gpt3/backbone.py +++ b/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 diff --git a/modelscope/models/nlp/gpt3/distributed_gpt3.py b/modelscope/models/nlp/gpt3/distributed_gpt3.py index 538631b6..c7917b31 100644 --- a/modelscope/models/nlp/gpt3/distributed_gpt3.py +++ b/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) diff --git a/modelscope/models/nlp/gpt3/text_generation.py b/modelscope/models/nlp/gpt3/text_generation.py index fbc82b8a..956e4553 100644 --- a/modelscope/models/nlp/gpt3/text_generation.py +++ b/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) diff --git a/modelscope/msdatasets/data_files/data_files_manager.py b/modelscope/msdatasets/data_files/data_files_manager.py index a84876f6..d5f43533 100644 --- a/modelscope/msdatasets/data_files/data_files_manager.py +++ b/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) diff --git a/modelscope/msdatasets/data_loader/data_loader.py b/modelscope/msdatasets/data_loader/data_loader.py index b1450c61..f29acc8f 100644 --- a/modelscope/msdatasets/data_loader/data_loader.py +++ b/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() diff --git a/modelscope/msdatasets/data_loader/data_loader_manager.py b/modelscope/msdatasets/data_loader/data_loader_manager.py index 5be32de1..0dec5d89 100644 --- a/modelscope/msdatasets/data_loader/data_loader_manager.py +++ b/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} .' diff --git a/modelscope/msdatasets/dataset_cls/custom_datasets/ocr_recognition_dataset.py b/modelscope/msdatasets/dataset_cls/custom_datasets/ocr_recognition_dataset.py index bfbb6eb3..07c6cdcc 100644 --- a/modelscope/msdatasets/dataset_cls/custom_datasets/ocr_recognition_dataset.py +++ b/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') diff --git a/modelscope/msdatasets/dataset_cls/dataset.py b/modelscope/msdatasets/dataset_cls/dataset.py index 9114285e..cebfcfba 100644 --- a/modelscope/msdatasets/dataset_cls/dataset.py +++ b/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): diff --git a/modelscope/msdatasets/download/dataset_builder.py b/modelscope/msdatasets/download/dataset_builder.py index 8ad5243a..796e3d83 100644 --- a/modelscope/msdatasets/download/dataset_builder.py +++ b/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 = {} diff --git a/modelscope/msdatasets/download/download_config.py b/modelscope/msdatasets/download/download_config.py index 4af656e0..11118f85 100644 --- a/modelscope/msdatasets/download/download_config.py +++ b/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 diff --git a/modelscope/msdatasets/meta/data_meta_manager.py b/modelscope/msdatasets/meta/data_meta_manager.py index 0fa74c37..3f1e6572 100644 --- a/modelscope/msdatasets/meta/data_meta_manager.py +++ b/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 diff --git a/modelscope/msdatasets/ms_dataset.py b/modelscope/msdatasets/ms_dataset.py index 912e061d..b9ff9971 100644 --- a/modelscope/msdatasets/ms_dataset.py +++ b/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 diff --git a/modelscope/msdatasets/utils/dataset_utils.py b/modelscope/msdatasets/utils/dataset_utils.py index dde044d5..dc7f9b38 100644 --- a/modelscope/msdatasets/utils/dataset_utils.py +++ b/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: diff --git a/modelscope/msdatasets/utils/oss_utils.py b/modelscope/msdatasets/utils/oss_utils.py index 44ee5446..d11b5035 100644 --- a/modelscope/msdatasets/utils/oss_utils.py +++ b/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 diff --git a/modelscope/pipelines/audio/asr_inference_pipeline.py b/modelscope/pipelines/audio/asr_inference_pipeline.py index b9c0bd03..823964e5 100644 --- a/modelscope/pipelines/audio/asr_inference_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/lm_infer_pipeline.py b/modelscope/pipelines/audio/lm_infer_pipeline.py index 75d835d6..e1524ebd 100644 --- a/modelscope/pipelines/audio/lm_infer_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/punctuation_processing_pipeline.py b/modelscope/pipelines/audio/punctuation_processing_pipeline.py index 3ab3481d..4e41e0c0 100644 --- a/modelscope/pipelines/audio/punctuation_processing_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/speaker_diarization_pipeline.py b/modelscope/pipelines/audio/speaker_diarization_pipeline.py index 71715ecd..a20cfcad 100644 --- a/modelscope/pipelines/audio/speaker_diarization_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/speaker_verification_pipeline.py b/modelscope/pipelines/audio/speaker_verification_pipeline.py index e576885a..c10f6a95 100644 --- a/modelscope/pipelines/audio/speaker_verification_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/timestamp_pipeline.py b/modelscope/pipelines/audio/timestamp_pipeline.py index 0968b359..17cf9545 100644 --- a/modelscope/pipelines/audio/timestamp_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/voice_activity_detection_pipeline.py b/modelscope/pipelines/audio/voice_activity_detection_pipeline.py index 0121b242..3e00454a 100644 --- a/modelscope/pipelines/audio/voice_activity_detection_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py index dd39453c..ca1431ea 100644 --- a/modelscope/pipelines/builder.py +++ b/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 diff --git a/modelscope/pipelines/cv/card_detection_pipeline.py b/modelscope/pipelines/cv/card_detection_pipeline.py index 75e7743a..5d7e87b0 100644 --- a/modelscope/pipelines/cv/card_detection_pipeline.py +++ b/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() diff --git a/modelscope/pipelines/cv/face_detection_pipeline.py b/modelscope/pipelines/cv/face_detection_pipeline.py index 58aefedf..115e8998 100644 --- a/modelscope/pipelines/cv/face_detection_pipeline.py +++ b/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() diff --git a/modelscope/pipelines/cv/ocr_recognition_pipeline.py b/modelscope/pipelines/cv/ocr_recognition_pipeline.py index e81e1ff6..602e23d4 100644 --- a/modelscope/pipelines/cv/ocr_recognition_pipeline.py +++ b/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): diff --git a/modelscope/pipelines/multi_modal/__init__.py b/modelscope/pipelines/multi_modal/__init__.py index b28e9a71..11c28fdd 100644 --- a/modelscope/pipelines/multi_modal/__init__.py +++ b/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'], diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py index 0c9fc5e8..8c5a2632 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/__init__.py +++ b/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 diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py index 6892877a..5f5d44f3 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/__init__.py +++ b/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'] } diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py index 49b4ef37..c4f8bda8 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/stable_diffusion/stable_diffusion_pipeline.py +++ b/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 = [] diff --git a/modelscope/pipelines/nlp/distributed_gpt3_pipeline.py b/modelscope/pipelines/nlp/distributed_gpt3_pipeline.py index 1738f2da..6962a09a 100644 --- a/modelscope/pipelines/nlp/distributed_gpt3_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/nlp/text_generation_pipeline.py b/modelscope/pipelines/nlp/text_generation_pipeline.py index d1aa5ff6..cfc3645d 100644 --- a/modelscope/pipelines/nlp/text_generation_pipeline.py +++ b/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], diff --git a/modelscope/preprocessors/multi_modal.py b/modelscope/preprocessors/multi_modal.py index faf796f4..82d44da8 100644 --- a/modelscope/preprocessors/multi_modal.py +++ b/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), diff --git a/modelscope/preprocessors/nlp/token_classification_preprocessor.py b/modelscope/preprocessors/nlp/token_classification_preprocessor.py index 4b4fee1f..b3ff9935 100644 --- a/modelscope/preprocessors/nlp/token_classification_preprocessor.py +++ b/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)] diff --git a/modelscope/trainers/default_config.py b/modelscope/trainers/default_config.py index bb272695..e4b08b4a 100644 --- a/modelscope/trainers/default_config.py +++ b/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 diff --git a/modelscope/trainers/hooks/checkpoint/checkpoint_hook.py b/modelscope/trainers/hooks/checkpoint/checkpoint_hook.py index 4b14a13f..b8be0682 100644 --- a/modelscope/trainers/hooks/checkpoint/checkpoint_hook.py +++ b/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 diff --git a/modelscope/trainers/hooks/distributed/deepspeed_hook.py b/modelscope/trainers/hooks/distributed/deepspeed_hook.py deleted file mode 100644 index 7dddc5d9..00000000 --- a/modelscope/trainers/hooks/distributed/deepspeed_hook.py +++ /dev/null @@ -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 diff --git a/modelscope/trainers/hooks/lr_scheduler_hook.py b/modelscope/trainers/hooks/lr_scheduler_hook.py index 51a8e858..facf5155 100644 --- a/modelscope/trainers/hooks/lr_scheduler_hook.py +++ b/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] diff --git a/modelscope/trainers/multi_modal/lora_diffusion/__init__.py b/modelscope/trainers/multi_modal/lora_diffusion/__init__.py new file mode 100644 index 00000000..311d2789 --- /dev/null +++ b/modelscope/trainers/multi_modal/lora_diffusion/__init__.py @@ -0,0 +1 @@ +from .lora_diffusion_trainer import LoraDiffusionTrainer diff --git a/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py b/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py new file mode 100644 index 00000000..40da164e --- /dev/null +++ b/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 diff --git a/modelscope/trainers/multi_modal/stable_diffusion/__init__.py b/modelscope/trainers/multi_modal/stable_diffusion/__init__.py new file mode 100644 index 00000000..a23f2880 --- /dev/null +++ b/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 diff --git a/modelscope/trainers/multi_modal/stable_diffusion/stable_diffusion_trainer.py b/modelscope/trainers/multi_modal/stable_diffusion/stable_diffusion_trainer.py new file mode 100644 index 00000000..68d7c689 --- /dev/null +++ b/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 diff --git a/modelscope/trainers/trainer.py b/modelscope/trainers/trainer.py index c980de04..fd0fafb8 100644 --- a/modelscope/trainers/trainer.py +++ b/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 diff --git a/modelscope/trainers/training_args.py b/modelscope/trainers/training_args.py index b7236163..aca3a886 100644 --- a/modelscope/trainers/training_args.py +++ b/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']) diff --git a/modelscope/tuners/control_sd_lora.py b/modelscope/tuners/control_sd_lora.py index 2585daa1..a2c53e24 100644 --- a/modelscope/tuners/control_sd_lora.py +++ b/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() diff --git a/modelscope/tuners/sd_lora.py b/modelscope/tuners/sd_lora.py index d740f7c0..feff05f4 100644 --- a/modelscope/tuners/sd_lora.py +++ b/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( diff --git a/modelscope/utils/typing.py b/modelscope/utils/input_output_typing.py similarity index 100% rename from modelscope/utils/typing.py rename to modelscope/utils/input_output_typing.py diff --git a/modelscope/utils/plugins.py b/modelscope/utils/plugins.py index 9d238e7d..e997f676 100644 --- a/modelscope/utils/plugins.py +++ b/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( diff --git a/modelscope/utils/pre_compile.py b/modelscope/utils/pre_compile.py new file mode 100644 index 00000000..65bde30a --- /dev/null +++ b/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() diff --git a/modelscope/utils/streaming_output.py b/modelscope/utils/streaming_output.py new file mode 100644 index 00000000..1a011a62 --- /dev/null +++ b/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 diff --git a/modelscope/version.py b/modelscope/version.py index 02464c29..32823e7b 100644 --- a/modelscope/version.py +++ b/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' diff --git a/requirements/audio/audio_asr.txt b/requirements/audio/audio_asr.txt index 7725a0dd..d136ce55 100644 --- a/requirements/audio/audio_asr.txt +++ b/requirements/audio/audio_asr.txt @@ -1,2 +1 @@ -easyasr>=0.0.2 -funasr>=0.5.0 +funasr>=0.6.0 diff --git a/requirements/audio/audio_kws.txt b/requirements/audio/audio_kws.txt index 4118f3ed..276a0a2f 100644 --- a/requirements/audio/audio_kws.txt +++ b/requirements/audio/audio_kws.txt @@ -1,7 +1,6 @@ kaldiio kwsbp>=0.0.6 matplotlib -numpy py_sound_connect>=0.1 scipy SoundFile>0.10 diff --git a/requirements/audio/audio_signal.txt b/requirements/audio/audio_signal.txt index 16a18e67..62944b71 100644 --- a/requirements/audio/audio_signal.txt +++ b/requirements/audio/audio_signal.txt @@ -2,7 +2,6 @@ hyperpyyaml librosa==0.9.2 MinDAEC mir_eval>=0.7 -numpy rotary_embedding_torch>=0.1.5 scipy SoundFile>0.10 diff --git a/requirements/cv.txt b/requirements/cv.txt index 0cec3659..c48d82e2 100644 --- a/requirements/cv.txt +++ b/requirements/cv.txt @@ -16,6 +16,7 @@ fastai>=1.0.51 ffmpeg>=1.4 ffmpeg-python>=0.2.0 ftfy +fvcore imageio>=2.9.0 imageio-ffmpeg>=0.4.2 imgaug>=0.4.0 @@ -44,16 +45,19 @@ pandas panopticapi plyfile>=0.7.4 psutil +pyclipper PyMCubes pytorch-lightning regex -scikit-image>=0.19.3 +# <0.20.0 for compatible python3.7 python3.8 +scikit-image>=0.19.3,<0.20.0 scikit-learn>=0.20.1 shapely shotdetect_scenedetect_lgss>=0.0.4 smplx tensorflow-estimator>=1.15.1 tf_slim +thop timm>=0.4.9 torchmetrics>=0.6.2 torchsummary>=1.5.1 diff --git a/requirements/framework.txt b/requirements/framework.txt index e763ae63..3ba659fc 100644 --- a/requirements/framework.txt +++ b/requirements/framework.txt @@ -7,8 +7,8 @@ gast>=0.2.2 # for python3.7 python3.8 compatible numpy<=1.22.0 oss2 -# for datasets compatible -pandas<=1.5.3 +# for datasets compatible and py37 py38 compatible +pandas<1.4.0 Pillow>=6.2.0 # pyarrow 9.0.0 introduced event_loop core dump pyarrow>=6.0.0,!=9.0.0 diff --git a/requirements/multi-modal.txt b/requirements/multi-modal.txt index 9d2c3448..fe2d45a1 100644 --- a/requirements/multi-modal.txt +++ b/requirements/multi-modal.txt @@ -1,10 +1,14 @@ accelerate -diffusers>=0.13.1,<0.15.0 +cloudpickle +decord>=0.6.0 +diffusers==0.15.0 +fairseq ftfy>=6.0.3 librosa==0.9.2 opencv-python pycocoevalcap>=1.2 pycocotools>=2.0.4 +pydot # compatible with taming-transformers-rom1504 pytorch_lightning<=1.7.7 rapidfuzz diff --git a/requirements/nlp.txt b/requirements/nlp.txt index 2e22ef8f..18f9a08f 100644 --- a/requirements/nlp.txt +++ b/requirements/nlp.txt @@ -1,4 +1,5 @@ boto3 +embeddings en_core_web_sm>=2.3.5 filelock ftfy @@ -18,6 +19,7 @@ scikit_learn sentencepiece seqeval spacy>=2.3.5 +stanza subword_nmt>=0.3.8 termcolor tokenizers diff --git a/requirements/tensorflow1x.txt b/requirements/tensorflow1x.txt index b139efe1..5d680652 100644 --- a/requirements/tensorflow1x.txt +++ b/requirements/tensorflow1x.txt @@ -1 +1 @@ -numpy==1.18.5 +numpy<1.20.0 diff --git a/tests/export/test_export_object_detection_damoyolo.py b/tests/export/test_export_object_detection_damoyolo.py index d7e51165..9a810e93 100644 --- a/tests/export/test_export_object_detection_damoyolo.py +++ b/tests/export/test_export_object_detection_damoyolo.py @@ -27,6 +27,15 @@ class TestExportObjectDetectionDamoyolo(unittest.TestCase): Exporter.from_model(model).export_onnx( input_shape=(1, 3, 640, 640), output_dir=self.tmp_dir) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_export_domain_specific_object_detection_damoyolo(self): + + model_id = 'damo/cv_tinynas_human-detection_damoyolo' + model = Model.from_pretrained(model_id) + with tempfile.TemporaryDirectory() as tmp_dir: + Exporter.from_model(model).export_onnx( + input_shape=(1, 3, 640, 640), output_dir=tmp_dir) + if __name__ == '__main__': unittest.main() diff --git a/tests/msdatasets/test_ms_dataset.py b/tests/msdatasets/test_ms_dataset.py index ddb84b45..83ffd3f8 100644 --- a/tests/msdatasets/test_ms_dataset.py +++ b/tests/msdatasets/test_ms_dataset.py @@ -241,8 +241,7 @@ class MsDatasetTest(unittest.TestCase): use_streaming=True) data_example = next(iter(dataset)) print(data_example) - assert isinstance(data_example['Noisy Image:FILE:Object'], - PngImageFile) + assert data_example.values() @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_to_ms_dataset(self): diff --git a/tests/pipelines/test_diffusers_stable_diffusion.py b/tests/pipelines/test_diffusers_stable_diffusion.py index eef677fc..57eae4a3 100644 --- a/tests/pipelines/test_diffusers_stable_diffusion.py +++ b/tests/pipelines/test_diffusers_stable_diffusion.py @@ -17,11 +17,11 @@ class DiffusersStableDiffusionTest(unittest.TestCase): test_input = 'a photo of an astronaut riding a horse on mars' - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run(self): diffusers_pipeline = pipeline(task=self.task, model=self.model_id) output = diffusers_pipeline({ - 'text': self.test_input, + 'prompt': self.test_input, 'height': 512, 'width': 512 }) diff --git a/tests/pipelines/test_efficient_diffusion_tuning.py b/tests/pipelines/test_efficient_diffusion_tuning.py index e33b2bf2..f1aa52de 100644 --- a/tests/pipelines/test_efficient_diffusion_tuning.py +++ b/tests/pipelines/test_efficient_diffusion_tuning.py @@ -13,7 +13,7 @@ class EfficientDiffusionTuningTest(unittest.TestCase): def setUp(self) -> None: self.task = Tasks.efficient_diffusion_tuning - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_efficient_diffusion_tuning_lora_run_pipeline(self): model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora' inputs = {'prompt': 'pale golden rod circle with old lace background'} diff --git a/tests/pipelines/test_gpt3_text_generation.py b/tests/pipelines/test_gpt3_text_generation.py index 1d938384..832295de 100644 --- a/tests/pipelines/test_gpt3_text_generation.py +++ b/tests/pipelines/test_gpt3_text_generation.py @@ -22,6 +22,13 @@ class TextGPT3GenerationTest(unittest.TestCase): pipe = pipeline(Tasks.text_generation, model=self.model_id_1_3B) print(pipe(self.input)) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_gpt3_1_3B_with_streaming(self): + pipe = pipeline(Tasks.text_generation, model=self.model_id_1_3B) + for output in pipe.stream(self.input, max_length=64): + print(output, end='\r') + print() + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_gpt3_2_7B(self): pipe = pipeline(Tasks.text_generation, model=self.model_id_2_7B) diff --git a/tests/pipelines/test_ocr_detection.py b/tests/pipelines/test_ocr_detection.py index 0ed2e59c..f5fc0c63 100644 --- a/tests/pipelines/test_ocr_detection.py +++ b/tests/pipelines/test_ocr_detection.py @@ -13,6 +13,7 @@ class OCRDetectionTest(unittest.TestCase): self.model_id = 'damo/cv_resnet18_ocr-detection-line-level_damo' self.model_id_vlpt = 'damo/cv_resnet50_ocr-detection-vlpt' self.model_id_db = 'damo/cv_resnet18_ocr-detection-db-line-level_damo' + self.model_id_db_nas = 'damo/cv_proxylessnas_ocr-detection-db-line-level_damo' self.test_image = 'data/test/images/ocr_detection.jpg' self.test_image_vlpt = 'data/test/images/ocr_detection_vlpt.jpg' self.task = Tasks.ocr_detection @@ -37,6 +38,15 @@ class OCRDetectionTest(unittest.TestCase): ocr_detection = pipeline(Tasks.ocr_detection, model=self.model_id_db) self.pipeline_inference(ocr_detection, self.test_image) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_dbnas_with_model_from_modelhub(self): + ocr_detection = pipeline( + Tasks.ocr_detection, + model=self.model_id_db_nas, + model_revision='v1.0.0', + ) + self.pipeline_inference(ocr_detection, self.test_image) + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_run_modelhub_default_model(self): ocr_detection = pipeline(Tasks.ocr_detection) diff --git a/tests/pipelines/test_ocr_recognition.py b/tests/pipelines/test_ocr_recognition.py index 94ee521f..27870b10 100644 --- a/tests/pipelines/test_ocr_recognition.py +++ b/tests/pipelines/test_ocr_recognition.py @@ -20,6 +20,21 @@ class OCRRecognitionTest(unittest.TestCase): result = pipeline(input_location) print('ocr recognition results: ', result) + def pipeline_inference_batch(self, pipeline: Pipeline, + input_location: str): + result = pipeline( + [input_location, input_location, input_location, input_location], + batch_size=4) + print('ocr recognition results: ', result) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_model_from_modelhub_batch(self): + ocr_recognition = pipeline( + Tasks.ocr_recognition, + model=self.model_id, + model_revision='v2.3.0') + self.pipeline_inference_batch(ocr_recognition, self.test_image) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_with_model_from_modelhub(self): ocr_recognition = pipeline( @@ -68,6 +83,14 @@ class OCRRecognitionTest(unittest.TestCase): model_revision='v2.2.2') self.pipeline_inference(ocr_recognition, self.test_image) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_model_from_modelhub_lightweightedge(self): + ocr_recognition = pipeline( + Tasks.ocr_recognition, + model='damo/cv_LightweightEdge_ocr-recognitoin-general_damo', + model_revision='v1.0.0') + self.pipeline_inference(ocr_recognition, self.test_image) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_model_from_modelhub_PILinput(self): ocr_recognition = pipeline( @@ -137,6 +160,15 @@ class OCRRecognitionTest(unittest.TestCase): device='cpu') self.pipeline_inference(ocr_recognition, self.test_image) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_model_from_modelhub_lightweightedge_cpu(self): + ocr_recognition = pipeline( + Tasks.ocr_recognition, + model='damo/cv_LightweightEdge_ocr-recognitoin-general_damo', + model_revision='v1.0.0', + device='cpu') + self.pipeline_inference(ocr_recognition, self.test_image) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_model_from_modelhub_PILinput_cpu(self): ocr_recognition = pipeline( diff --git a/tests/pipelines/test_text_generation.py b/tests/pipelines/test_text_generation.py index 378b1bbc..687a38df 100644 --- a/tests/pipelines/test_text_generation.py +++ b/tests/pipelines/test_text_generation.py @@ -8,6 +8,7 @@ from modelscope.pipelines import pipeline from modelscope.pipelines.nlp import TextGenerationPipeline from modelscope.preprocessors import TextGenerationTransformersPreprocessor from modelscope.utils.constant import Tasks +from modelscope.utils.streaming_output import StreamingOutputMixin from modelscope.utils.test_utils import test_level @@ -61,6 +62,20 @@ class TextGenerationTest(unittest.TestCase): task=Tasks.text_generation, model=model_id, **init_kwargs) print(pipeline_ins(input, **run_kwargs)) + def run_streaming_pipeline_with_model_id(self, + model_id, + input, + init_kwargs={}, + run_kwargs={}): + pipeline_ins = pipeline( + task=Tasks.text_generation, model=model_id, **init_kwargs) + + # set stream inputs + assert isinstance(pipeline_ins, StreamingOutputMixin) + for output in pipeline_ins.stream(input, **run_kwargs): + print(output, end='\r') + print() + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_palm_zh_base_with_model_name(self): self.run_pipeline_with_model_id(self.palm_model_id_zh_base, @@ -123,6 +138,23 @@ class TextGenerationTest(unittest.TestCase): [self.gpt3_input, self.gpt3_input[:10], self.gpt3_input[10:]], run_kwargs={'batch_size': 2}) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_gpt_base_with_model_name_with_streaming(self): + self.run_streaming_pipeline_with_model_id( + self.gpt3_base_model_id, + self.gpt3_input, + run_kwargs={'max_length': 64}) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_gpt_base_with_model_name_with_streaming_batch(self): + self.run_streaming_pipeline_with_model_id( + self.gpt3_base_model_id, + [self.gpt3_input, self.gpt3_input[:10], self.gpt3_input[10:]], + run_kwargs={ + 'batch_size': 2, + 'max_length': 32 + }) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_gpt_base_with_model_name_batch_iter(self): self.run_pipeline_with_model_id( diff --git a/tests/run_config.yaml b/tests/run_config.yaml index ba678468..cfa08e4a 100644 --- a/tests/run_config.yaml +++ b/tests/run_config.yaml @@ -60,48 +60,19 @@ isolated: # test cases that may require excessive anmount of GPU memory or run - test_video_deinterlace.py - test_image_inpainting_sdv2.py - test_bad_image_detecting.py - - test_image_portrait_stylization_trainer.py - test_controllable_image_generation.py - test_image_colorization_trainer.py envs: default: # default env, case not in other env will in default, pytorch. dependencies: # requirement packages，pip install before test case run. - - numpy>=1.20,<=1.21.0 + - numpy>=1.20,<=1.22.0 - protobuf<4,>=3.20.2 + tensorflow1x: # cases excuted tensorflow1.x framework. requirements: # requirements files run before test case run. - tensorflow1x.txt dependencies: # requirement packages，pip install before test case run. - - numpy==1.18.5 + - numpy<1.20.0 tests: - - test_text_to_speech.py - - test_csanmt_translation.py - - test_translation_trainer.py - - test_translation_evaluation_trainer.py - - test_ocr_detection.py - - test_automatic_speech_recognition.py - - test_image_matting.py - - test_person_image_cartoon.py - - test_skin_retouching.py - - test_image_style_transfer.py - test_image_portrait_stylization_trainer.py - - test_language_identification.py - - test_language_guided_video_summarization_trainer.py - - test_motion_generation.py - - test_universal_matting.py - - test_dialog_modeling.py - - test_trainer.py - - test_abnormal_object_detection.py - - test_image_face_fusion.py - - test_ocr_detection_db_trainer.py - - test_language_guided_video_summarization.py - - test_interactive_translation_pipeline.py - - test_image_defrcn_fewshot_trainer.py - - test_automatic_post_editing.py - - test_human_reconstruction.py - - test_nerf_recon_acc_trainer.py - - test_nerf_recon_acc.py - - test_speech_signal_process.py - - test_tensorboard_hook.py - - test_efficient_diffusion_tuning_trainer.py diff --git a/tests/trainers/cli/test_cli.py b/tests/trainers/cli/test_cli.py index b9fb7539..55a12b88 100644 --- a/tests/trainers/cli/test_cli.py +++ b/tests/trainers/cli/test_cli.py @@ -1,4 +1,5 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import os import unittest import json @@ -21,7 +22,7 @@ class TestCli(unittest.TestCase): 'sentence2': 'sentence2', 'label': 'label', }, - 'split': 0.8, + 'usage': 0.8, }, { 'dataset': { 'dataset_name': 'glue', @@ -33,11 +34,15 @@ class TestCli(unittest.TestCase): 'hypothesis': 'sentence2', 'label': 'label', }, - 'split': 'val', + 'usage': 'val', }] with open('./dataset.json', 'w') as f: json.dump(content, f) + def tearDown(self) -> None: + if os.path.exists('./dataset.json'): + os.remove('./dataset.json') + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_merge_dataset_from_file(self): dataset = MsDataset.load('clue', subset_name='cmnli', split='train') diff --git a/tests/trainers/test_efficient_diffusion_tuning_trainer.py b/tests/trainers/test_efficient_diffusion_tuning_trainer.py index 2484e24d..a19bf21d 100644 --- a/tests/trainers/test_efficient_diffusion_tuning_trainer.py +++ b/tests/trainers/test_efficient_diffusion_tuning_trainer.py @@ -39,7 +39,7 @@ class TestEfficientDiffusionTuningTrainer(unittest.TestCase): shutil.rmtree(self.tmp_dir) super().tearDown() - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_efficient_diffusion_tuning_lora_train(self): model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora' @@ -67,7 +67,7 @@ class TestEfficientDiffusionTuningTrainer(unittest.TestCase): for i in range(self.max_epochs): self.assertIn(f'epoch_{i+1}.pth', results_files) - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_efficient_diffusion_tuning_lora_eval(self): model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora' @@ -87,7 +87,7 @@ class TestEfficientDiffusionTuningTrainer(unittest.TestCase): result = trainer.evaluate() print(f'Efficient-diffusion-tuning-lora eval output: {result}.') - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_efficient_diffusion_tuning_control_lora_train(self): model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora' @@ -116,7 +116,7 @@ class TestEfficientDiffusionTuningTrainer(unittest.TestCase): for i in range(self.max_epochs): self.assertIn(f'epoch_{i+1}.pth', results_files) - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_efficient_diffusion_tuning_control_lora_eval(self): model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora' diff --git a/tests/trainers/test_lora_diffusion_trainer.py b/tests/trainers/test_lora_diffusion_trainer.py new file mode 100644 index 00000000..2ba89665 --- /dev/null +++ b/tests/trainers/test_lora_diffusion_trainer.py @@ -0,0 +1,89 @@ +# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved. +import os +import shutil +import tempfile +import unittest + +from modelscope.metainfo import Trainers +from modelscope.msdatasets import MsDataset +from modelscope.trainers import build_trainer +from modelscope.utils.constant import DownloadMode +from modelscope.utils.test_utils import test_level + + +class TestLoraDiffusionTrainer(unittest.TestCase): + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + + self.train_dataset = MsDataset.load( + 'buptwq/lora-stable-diffusion-finetune', + split='train', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + self.eval_dataset = MsDataset.load( + 'buptwq/lora-stable-diffusion-finetune', + split='validation', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + + self.max_epochs = 5 + + self.tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(self.tmp_dir): + os.makedirs(self.tmp_dir) + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + super().tearDown() + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_lora_diffusion_train(self): + model_id = 'AI-ModelScope/stable-diffusion-v1-5' + model_revision = 'v1.0.6' + + def cfg_modify_fn(cfg): + cfg.train.max_epochs = self.max_epochs + cfg.train.lr_scheduler = { + 'type': 'LambdaLR', + 'lr_lambda': lambda _: 1, + 'last_epoch': -1 + } + cfg.train.optimizer.lr = 1e-4 + return cfg + + kwargs = dict( + model=model_id, + model_revision=model_revision, + work_dir=self.tmp_dir, + train_dataset=self.train_dataset, + eval_dataset=self.eval_dataset, + cfg_modify_fn=cfg_modify_fn) + + trainer = build_trainer( + name=Trainers.lora_diffusion, default_args=kwargs) + trainer.train() + result = trainer.evaluate() + print(f'Lora-diffusion train output: {result}.') + + results_files = os.listdir(self.tmp_dir) + self.assertIn(f'{trainer.timestamp}.log.json', results_files) + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_lora_diffusion_eval(self): + model_id = 'AI-ModelScope/stable-diffusion-v1-5' + model_revision = 'v1.0.6' + + kwargs = dict( + model=model_id, + model_revision=model_revision, + work_dir=self.tmp_dir, + train_dataset=None, + eval_dataset=self.eval_dataset) + + trainer = build_trainer( + name=Trainers.lora_diffusion, default_args=kwargs) + result = trainer.evaluate() + print(f'Lora-diffusion eval output: {result}.') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/trainers/test_stable_diffusion_trainer.py b/tests/trainers/test_stable_diffusion_trainer.py new file mode 100644 index 00000000..c3dbcdee --- /dev/null +++ b/tests/trainers/test_stable_diffusion_trainer.py @@ -0,0 +1,89 @@ +# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved. +import os +import shutil +import tempfile +import unittest + +from modelscope.metainfo import Trainers +from modelscope.msdatasets import MsDataset +from modelscope.trainers import build_trainer +from modelscope.utils.constant import DownloadMode +from modelscope.utils.test_utils import test_level + + +class TestStableDiffusionTrainer(unittest.TestCase): + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + + self.train_dataset = MsDataset.load( + 'buptwq/lora-stable-diffusion-finetune', + split='train', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + self.eval_dataset = MsDataset.load( + 'buptwq/lora-stable-diffusion-finetune', + split='validation', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + + self.max_epochs = 5 + + self.tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(self.tmp_dir): + os.makedirs(self.tmp_dir) + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + super().tearDown() + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_stable_diffusion_train(self): + model_id = 'AI-ModelScope/stable-diffusion-v1-5' + model_revision = 'v1.0.7' + + def cfg_modify_fn(cfg): + cfg.train.max_epochs = self.max_epochs + cfg.train.lr_scheduler = { + 'type': 'LambdaLR', + 'lr_lambda': lambda _: 1, + 'last_epoch': -1 + } + cfg.train.optimizer.lr = 1e-4 + return cfg + + kwargs = dict( + model=model_id, + model_revision=model_revision, + work_dir=self.tmp_dir, + train_dataset=self.train_dataset, + eval_dataset=self.eval_dataset, + cfg_modify_fn=cfg_modify_fn) + + trainer = build_trainer( + name=Trainers.stable_diffusion, default_args=kwargs) + trainer.train() + result = trainer.evaluate() + print(f'Stable-diffusion train output: {result}.') + + results_files = os.listdir(self.tmp_dir) + self.assertIn(f'{trainer.timestamp}.log.json', results_files) + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_stable_diffusion_eval(self): + model_id = 'AI-ModelScope/stable-diffusion-v1-5' + model_revision = 'v1.0.7' + + kwargs = dict( + model=model_id, + model_revision=model_revision, + work_dir=self.tmp_dir, + train_dataset=None, + eval_dataset=self.eval_dataset) + + trainer = build_trainer( + name=Trainers.stable_diffusion, default_args=kwargs) + result = trainer.evaluate() + print(f'Stable-diffusion eval output: {result}.') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/utils/test_ast.py b/tests/utils/test_ast.py index 78917ccd..288c076a 100644 --- a/tests/utils/test_ast.py +++ b/tests/utils/test_ast.py @@ -47,7 +47,7 @@ class AstScaningTest(unittest.TestCase): self.assertIsInstance(decorators, list) self.assertListEqual( list(set(imports.keys()) - set(['torch', 'os'])), []) - self.assertEqual(len(from_imports.keys()), 10) + self.assertEqual(len(from_imports.keys()), 11) self.assertTrue(from_imports['modelscope.metainfo'] is not None) self.assertEqual(from_imports['modelscope.metainfo'], ['Pipelines']) self.assertEqual(