Dockerfile.local

ARG SERVERBASE=21.06-py3

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

ENV FORCE_CUDA=1
ARG LIB_WITH_CUDA=ON
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.8/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.8/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.8/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.8/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 pushd pytorch_geometric && pip install -e . && popd

RUN pushd pytorch_geometric/examples/jit &&\
    echo "torch.jit.save(model, 'gat_test.pt')" >> gat.py &&\
    python gat.py >> train.log &&\
    mv gat_test.pt /workspace/model.pt

COPY gat_local_test.py /workspace/