44 lines
1.4 KiB
Bash
44 lines
1.4 KiB
Bash
#!/bin/bash
|
|
|
|
LIBTORCH_VERSION="${LIBTORCH_VER:-1.10.0}"
|
|
CUDA_VERSION="${CUDA_VER:-11.1}"
|
|
|
|
if [ "${CUDA_VERSION}"=="cpu" ]; then
|
|
CU_VERSION="cpu"
|
|
else
|
|
CU_VERSION="cu${CUDA_VERSION//./}"
|
|
fi
|
|
|
|
# Libtorch paths
|
|
LIBTORCH_PATH="/usr/local/lib"
|
|
GOTCH_LIBTORCH="$LIBRARY_PATH/libtorch"
|
|
|
|
LIBRARY_PATH="$LIBRARY_PATH:$GOTCH_LIBTORCH/lib"
|
|
CPATH="$CPATH:$GOTCH_LIBTORCH/lib:$GOTCH_LIBTORCH/include:$GOTCH_LIBTORCH/include/torch/csrc/api/include"
|
|
if [[ -z "${CUDA_VERSION}"=="cpu" ]]; then
|
|
LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$GOTCH_LIBTORCH/lib"
|
|
else
|
|
LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$GOTCH_LIBTORCH/lib:/usr/lib64-nvidia:/usr/local/cuda-${CUDA_VERSION}/lib64"
|
|
fi
|
|
|
|
# Update current shell environment variables for newly installed Libtorch
|
|
export GOTCH_LIBTORCH=$GOTCH_LIBTORCH
|
|
export LIBRARY_PATH=$LIBRARY_PATH
|
|
export CPATH=$CPATH
|
|
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH
|
|
|
|
# Install Libtorch
|
|
#=================
|
|
LIBTORCH_ZIP="libtorch-cxx11-abi-shared-with-deps-${LIBTORCH_VERSION}%2B${CU_VERSION}.zip"
|
|
LIBTORCH_URL="https://download.pytorch.org/libtorch/${CU_VERSION}/${LIBTORCH_ZIP}"
|
|
echo $LIBTORCH_URL
|
|
sudo wget -q --show-progress --progress=bar:force:noscroll -O "/tmp/$LIBTORCH_ZIP" "$LIBTORCH_URL"
|
|
sudo rm -rf $GOTCH_LIBTORCH # delete old libtorch if exisitng
|
|
sudo unzip "/tmp/$LIBTORCH_ZIP" -d $LIBTORCH_PATH
|
|
sudo rm "/tmp/$LIBTORCH_ZIP"
|
|
|
|
sudo ldconfig
|
|
|
|
# refresh environment for all next opening shells.
|
|
exec "$BASH"
|