Dockerfile.pyg

ARG SERVERBASE=20.09-py3

FROM nvcr.io/nvidia/pytorch:${SERVERBASE}

ENV FORCE_CUDA=1
ARG LIB_WITH_CUDA=ON
ARG NPROC=4

RUN git clone https://github.com/rusty1s/pytorch_cluster.git
RUN pushd pytorch_cluster &&\
    pip install . &&\
    mkdir build && pushd build &&\
    cmake -DCMAKE_PREFIX_PATH=/opt/conda/lib/python3.6/site-packages/torch -DWITH_CUDA=${LIB_WITH_CUDA} .. &&\
    make -j ${NPROC} && mv *.so /workspace/ && popd &&\
    popd

RUN git clone https://github.com/rusty1s/pytorch_scatter.git
RUN pushd pytorch_scatter &&\
    pip install . &&\
    mkdir build && pushd build &&\
    cmake -DCMAKE_PREFIX_PATH=/opt/conda/lib/python3.6/site-packages/torch -DWITH_CUDA=${LIB_WITH_CUDA} .. &&\
    make -j ${NPROC} && mv *.so /workspace/ && popd &&\
    popd

RUN git clone https://github.com/rusty1s/pytorch_spline_conv.git
RUN pushd pytorch_spline_conv &&\
    pip install . &&\
    mkdir build && pushd build &&\
    cmake -DCMAKE_PREFIX_PATH=/opt/conda/lib/python3.6/site-packages/torch -DWITH_CUDA=${LIB_WITH_CUDA} .. &&\
    make -j ${NPROC} && mv *.so /workspace/ && popd &&\
    popd

RUN git clone https://github.com/rusty1s/pytorch_sparse.git
RUN pushd pytorch_sparse &&\
    pip install . &&\
    mkdir build && pushd build &&\
    cmake -DCMAKE_PREFIX_PATH=/opt/conda/lib/python3.6/site-packages/torch -DWITH_CUDA=${LIB_WITH_CUDA} .. &&\
    make -j ${NPROC} && mv *.so /workspace/ && popd &&\
    popd

RUN git clone https://github.com/rusty1s/pytorch_geometric.git

RUN pip install llvmlite --ignore-installed
RUN pushd pytorch_geometric && pip install -e . && popd