pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
ARG PYTHON_VERSION=3.9
16
ARG PYTORCH_VERSION=1.13
17
ARG CUDA_VERSION=11.3.1
18

19
FROM pytorchlightning/pytorch_lightning:base-cuda-py${PYTHON_VERSION}-torch${PYTORCH_VERSION}-cuda${CUDA_VERSION}
20

21
LABEL maintainer="Lightning-AI <https://github.com/Lightning-AI>"
22

23
ARG LIGHTNING_VERSION=""
24

25
COPY ./ /home/pytorch-lightning/
26

27
# install dependencies
28
RUN \
29
    cd /home && \
30
    mv pytorch-lightning/_notebooks notebooks && \
31
    mv pytorch-lightning/examples . && \
32
    # replace by specific version if asked
33
    if [ ! -z "$LIGHTNING_VERSION" ] ; then \
34
        rm -rf pytorch-lightning ; \
35
        wget https://github.com/Lightning-AI/pytorch-lightning/archive/${LIGHTNING_VERSION}.zip --progress=bar:force:noscroll ; \
36
        unzip ${LIGHTNING_VERSION}.zip ; \
37
        mv pytorch-lightning-*/ pytorch-lightning ; \
38
        rm *.zip ; \
39
    fi && \
40
    # otherwise there is collision with folder name and pkg name on Pypi
41
    cd pytorch-lightning && \
42
    pip install setuptools==59.5.0 && \
43
    PACKAGE_NAME=lightning pip install '.[extra,loggers,strategies]' --no-cache-dir && \
44
    PACKAGE_NAME=pytorch pip install '.[extra,loggers,strategies]' --no-cache-dir && \
45
    cd .. && \
46
    rm -rf pytorch-lightning
47

48
RUN python --version && \
49
    pip --version && \
50
    pip list && \
51
    python -c "import pytorch_lightning as pl; print(pl.__version__)"
52

53
# CMD ["/bin/bash"]
54