mmaction2 の onnxエクスポート試行

git clone https://github.com/open-mmlab/mmaction2.git && cd mmaction2
git checkout 7c94243542985db813bb9021f97c95b59d136e52
mkdir -p data
mv docker/Dockerfile docker/Dockerfile_org

cat << 'EOT' > docker/Dockerfile
ARG PYTORCH="1.10.0"
ARG CUDA="11.3"
ARG CUDNN="8"

FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel

ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0+PTX 8.6"
ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
ENV CMAKE_PREFIX_PATH="$(dirname $(which conda))/../"
ARG PYTORCH="1.10.0"
ARG MMCUDA="113"
ARG MMCVFULL="1.4.4"
ARG USERNAME=user
ARG WKDIR=/home/${USERNAME}
WORKDIR ${WKDIR}

# Install dependencies
RUN apt-get update && apt-get install -y \
        git ninja-build libglib2.0-0 wget \
        libsm6 libxrender-dev libxext6 ffmpeg \
        sudo \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Install mmcv-full
RUN pip install pip --upgrade \
    && pip install onnx==1.10.2 \
    && pip install onnxruntime==1.10.0 \
    && pip install onnxruntime-extensions==0.4.2 \
    && pip install onnxoptimizer==0.2.6 \
    && pip install onnx-simplifier==0.3.6 \
    && pip install cython==0.29.27 --no-cache-dir \
    && pip install mmcv-full==${MMCVFULL} \
        -f https://download.openmmlab.com/mmcv/dist/cu${MMCUDA}/torch${PYTORCH}/index.html \
    && pip cache purge

# Install MMAction2
RUN conda clean --all
ENV FORCE_CUDA="1"
COPY . .
RUN pip install --no-cache-dir -e .

# Create a user who can sudo in the Docker container
RUN echo "root:root" | chpasswd \
    && adduser --disabled-password --gecos "" "${USERNAME}" \
    && echo "${USERNAME}:${USERNAME}" | chpasswd \
    && echo "%${USERNAME}    ALL=(ALL)   NOPASSWD:    ALL" >> /etc/sudoers.d/${USERNAME} \
    && chmod 0440 /etc/sudoers.d/${USERNAME}
USER ${USERNAME}
RUN sudo chown ${USERNAME}:${USERNAME} ${WKDIR} \
    && echo "cd ${WKDIR}/workdir" >> ${WKDIR}/.bashrc
EOT

https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html

docker build -t mmaction2test -f docker/Dockerfile .

docker run --rm -it --gpus all \
-v `pwd`:/home/user/workdir \
mmaction2test:latest

cat /etc/issue
Ubuntu 18.04.5 LTS \n \l

- Recognizers
 - TSN: $batch $clip $channel $height $width  1 1 3 224 224
 - I3D: $batch $clip $channel $time $height $width  1 1 3 32 224 224

python tools/deployment/pytorch2onnx.py \
${CONFIG_FILE} \
${CHECKPOINT_FILE} \
--shape ${SHAPE} \
--verify \
--output-file ${OUTPUT_FILE} \
--opset-version 11

========================================
- Localizers
 - BSN etc: Default  1 1 3 224 224

python tools/deployment/pytorch2onnx.py \
${CONFIG_FILE} \
${CHECKPOINT_FILE} \
--shape ${SHAPE} \
--verify \
--output-file ${OUTPUT_FILE} \
--is-localizer \
--opset-version 11

# PoseC3D - FineGYM
## config
wget https://github.com/open-mmlab/mmaction2/raw/master/configs/skeleton/posec3d/slowonly_r50_u48_240e_gym_keypoint.py
wget https://github.com/open-mmlab/mmaction2/raw/master/configs/skeleton/posec3d/slowonly_r50_u48_240e_gym_limb.py
## checkpoint
wget https://download.openmmlab.com/mmaction/skeleton/posec3d/slowonly_r50_u48_240e_gym_keypoint/slowonly_r50_u48_240e_gym_keypoint-b07a98a0.pth
wget https://download.openmmlab.com/mmaction/skeleton/posec3d/slowonly_r50_u48_240e_gym_limb/slowonly_r50_u48_240e_gym_limb-c0d7b482.pth

MODEL=slowonly_r50_u48_240e_gym_keypoint
H=224
W=224
python tools/deployment/pytorch2onnx.py \
${MODEL}.py \
${MODEL}-b07a98a0.pth \
--shape 1 1 17 32 ${H} ${W} \
--verify \
--output-file ${MODEL}_${H}x${W}.onnx \
--opset-version 11

python -m onnxsim ${MODEL}_${H}x${W}.onnx ${MODEL}_${H}x${W}.onnx

docker pull pinto0309/mmaction2_exec_env