Básicamente, he conseguido montar un entorno con torch y torchvision que funciona. La limitación es que por narices como tenemos jetpack 5.1.2 tiene que ser python3.8 y pytorch 2.1.0 con torchvision 0.16.1. Los paquetes de torch disponibles publicados por nvidia para jetson están aquí: https://developer.nvidia.com/w/compute/redist/jp/. Así he montado yo el entorno paso a paso (luego he instalado otras cosillas que necesitaba yo para mmsegmentation y he tenido que hacer alguna otra ñapa, pero específica de mi entorno):
conda create -n image-mmseg python=3.8 -y
conda activate image-mmseg
wget https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
pip install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
git clone --branch v0.16.1 https://github.com/pytorch/vision torchvision
export BUILD_VERSION=0.16.1
cd torchvision/
python3 setup.py install
pip install "pillow<11" requests
otra cosa que he descubierto es que hay una imagen de docker en la jetson (l4t-pytorch) que tiene pytorch 2.0.0 instalado y funcionando:
roboticslab@jetson-agx:~$ docker image ls
REPOSITORY TAG IMAGE ID CREATED SIZE
nvcr.io/nvidia/l4t-pytorch r35.2.1-pth2.0-py3 853b58c1dce6 3 years ago 11.7GB
roboticslab@jetson-agx:~$ docker run -it --rm --runtime nvidia --network host nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3
root@jetson-agx:/# python3 -c "import torch, torchvision; from PIL import Image; print(torch.__version__, torchvision.__version__); print(torch.cuda.is_available())"
2.0.0a0+ec3941ad.nv23.02 0.14.1a0+5e8e2f1
True
root@jetson-agx:/#
otras versiones de esa imagen de docker están disponibles aquí: https://catalog.ngc.nvidia.com/orgs/nvidia/-/containers/l4t-pytorch/-?_lr=1, pero dejaron de sacar nuevas hace 2 o 3 años (creo que para jetpack 6 sí que hay versiones actualizadas pero habría que reinstalar el cacharro entero)
para instalar mmsegmentation:
pip install numpy openmim==0.3.9 onnxruntime onnx ftfy regex
pip install lightning --no-deps
pip install lightning-utilities torchmetrics fsspec PyYAML tqdm packaging typing-extensions
pip install pytorch-lightning --no-deps
mim install mmengine==0.10.6 mmcv==2.1.0 mmsegmentation==1.2.2 mmdet==3.3.0
para instalar PerceptionMetrics:
git clone https://github.com/JdeRobot/PerceptionMetrics.git
eliminar los números de versión de los paquetes del pyproject.toml de PerceptionMetrics y:
cd PerceptionMetrics
pip install -e .
para evitar errores con torch distributed y mmsegmentation, modificamos los ficheros: