install_pytorch.sh
The install_pytorch.sh
script is designed to automate the installation of PyTorch in a Docker container, offering flexibility in the installation method.
It supports installing PyTorch from PyPI (Python Package Index) or directly from source, with additional handling for different C++ ABI (Application Binary Interface) settings.
The script adjusts to the host system's operating environment, specifically tailored for ubuntu
and centos
systems.
Here's a detailed breakdown of its components and functionalities:
Shebang and Bash Options
#!/bin/bash
: Declares the script to be run with Bash.set -ex
:-e
option makes the script exit immediately if a command exits with a non-zero status.-x
option prints each command that will be executed, aiding in debugging.
Variables
TORCH_VERSION
: Specifies the version of PyTorch to be installed.SYSTEM_ID
: Extracts the system identifier from/etc/os-release
, determining whether the operating system isubuntu
orcentos
.
Functions
prepare_environment
prepare_environment
Adjusts the environment before the PyTorch installation. For ubuntu
, it installs ninja-build
, a small build system with a focus on speed, and cleans up afterwards.
For centos
, it performs a similar installation and offers an option to temporarily disable the devtoolset
environment, which could be necessary to avoid conflicts or issues during the build process.
restore_environment
restore_environment
Reverts any changes made by prepare_environment
to the environment. Specifically, for centos
, it restores the devtoolset
environment if it was previously disabled.
install_from_source
install_from_source
Handles the installation of PyTorch from source.
It first prepares the environment, sets the _GLIBCXX_USE_CXX11_ABI
variable based on the function's argument to control the C++ ABI compatibility, and sets TORCH_CUDA_ARCH_LIST
to specify CUDA architectures for PyTorch to support.
It then proceeds to clone the PyTorch repository, install its Python dependencies, and build it from source. After installation, it cleans up and restores any environment changes.
install_from_pypi
install_from_pypi
Simplifies the PyTorch installation by using pip3
to install directly from PyPI, requiring only the version number.
Script Logic
The script uses a case statement to parse the first input argument ($1
), supporting four modes:
"skip"
: Does nothing, allowing the script to complete without installing PyTorch. Useful for containers where PyTorch is not needed or will be installed differently."pypi"
: Callsinstall_from_pypi
to install PyTorch from PyPI."src_cxx11_abi"
: Installs PyTorch from source with C++11 ABI compatibility enabled."src_non_cxx11_abi"
: Installs PyTorch from source without C++11 ABI compatibility.
If an incorrect argument is provided, it prints an error message and exits.
Summary
The install_pytorch.sh
script is a comprehensive tool for setting up PyTorch in Docker containers, accommodating different environments and installation preferences.
It demonstrates careful handling of system-specific dependencies, conditional logic based on system characteristics, and flexibility in managing complex software installations.
The ability to switch between PyTorch installation methods based on the ABI compatibility requirement is particularly useful for developers working in environments with specific C++ standards or in mixed ABI settings.
Last updated