Compile once and run on other GPU(different computational capacities)

[JiaMingLin]

Hi,

Thanks for this nice implementation, it brings convenient to our project.

I have slightly modify the setup.py script to compile for multiple GPU_CCs. And I belief it would benefit for someone.
In our case, we develop the prototype on local computer with desktop GPU, e.g. GTX1060, with small batch size. Then for tuning the model, we switch to datacenter where leverages Containers for fast deploy, and they offent suggest users not to compile code in the container since security issues. Using the original setup.py script, we have to re-compile for different GPU CCs. The following is my modification, I have tested for RTX 2070(CC=7.5), GTX1080(CC=6.1), V100(CC=7.0). And last, I provide my Dockerfile for quick reproduce.

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
from os.path import join

CPU_ONLY = False
project_root = 'Correlation_Module'

source_files = ['correlation.cpp', 'correlation_sampler.cpp']

cxx_args = ['-std=c++14', '-fopenmp']
nvcc_args = [
    '-gencode', 'arch=compute_50,code=sm_50',
    '-gencode', 'arch=compute_52,code=sm_52',
    '-gencode', 'arch=compute_60,code=sm_60',
    '-gencode', 'arch=compute_61,code=sm_61',
    '-gencode', 'arch=compute_70,code=sm_70',
    '-gencode', 'arch=compute_70,code=compute_70',
    '-gencode', 'arch=compute_75,code=compute_75',
    '-gencode', 'arch=compute_75,code=sm_75'
]

with open("README.md", "r") as fh:
    long_description = fh.read()


def launch_setup():
    if CPU_ONLY:
        Extension = CppExtension
        macro = []
    else:
        Extension = CUDAExtension
        source_files.append('correlation_cuda_kernel.cu')
        macro = [("USE_CUDA", None)]

    sources = [join(project_root, file) for file in source_files]

    setup(
        name='spatial_correlation_sampler',
        version="0.3.0",
        author="Clément Pinard",
        author_email="clement.pinard@ensta-paristech.fr",
        description="Correlation module for pytorch",
        long_description=long_description,
        long_description_content_type="text/markdown",
        url="https://github.com/ClementPinard/Pytorch-Correlation-extension",
        install_requires=['torch>=1.1', 'numpy'],
        ext_modules=[
            Extension('spatial_correlation_sampler_backend',
                      sources,
                      define_macros=macro,
                      extra_compile_args={'cxx': ['-fopenmp'], 'nvcc':nvcc_args},
                      extra_link_args=['-lgomp'])
        ],
        package_dir={'': project_root},
        packages=['spatial_correlation_sampler'],
        cmdclass={
            'build_ext': BuildExtension
        },
        classifiers=[
            "Programming Language :: Python :: 3",
            "License :: OSI Approved :: MIT License",
            "Operating System :: POSIX :: Linux",
            "Intended Audience :: Science/Research",
            "Topic :: Scientific/Engineering :: Artificial Intelligence"
        ])


if __name__ == '__main__':
    launch_setup()

Dockerfile

FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

RUN apt update -y

# install python
RUN apt-get install software-properties-common -y && \
    add-apt-repository ppa:deadsnakes/ppa -y && \
    apt-get update && \
    apt-get install python3.6 -y && \
    ln -s /usr/bin/python3 /usr/bin/python 

# install pip
RUN apt install python3-pip -y && \
    ln -s /usr/bin/pip3 /usr/bin/pip && \
    python -m pip install -U --force-reinstall pip

# install dependencies
RUN python -m pip install requests tqdm && \
    pip install torch==1.8.1+cu101 torchvision==0.9.1+cu101 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html && \
    apt install git vim -y

# install spatial correlation sampler
RUN git clone https://github.com/JiaMingLin/Pytorch-Correlation-extension.git && \
    cd Pytorch-Correlation-extension && \
    python setup.py install && \
    cd ..

[ClementPinard]

Hello, thank you for this nice addition !

Care to make a small PR ? That would add you as contributors of this project.

Slight comment about the code is that it would be nice to have a switch to add all other args or not, to avoid unnecessary building time and size when only compiling for desktop.

Cheers !