Pytorch support

Hi, I am exciting to see the announce and release of Jetson Nano. However, I could only get Tensorflow 1.13 working on Python 3.6. Could you provide Pytorch installation python wheel or installation guide?

It seems that Nvidia has some numbers on Pytorch on Nvidia Jetson Nano: https://devblogs.nvidia.com/jetson-nano-ai-computing/, is this test done on the real device, or you guys just provide some estimation?

you should have meant python 3.6 rather than pytorch 3.6 in the statement

I have a Python 2.7 wheel built of PyTorch master for Nano, Xavier, TX2, and JetPack 4.2. Would that be useful to you, or are you looking for a different Python version?

Awesome, can I have the link? Also, could you provide the wheel for python 3.6? Since the default python in jet pack 4.2 has python 3.6, which has tensorflow 1.13 built wheel.

It would be great if there was a way to get PyTorch V1.0 running with Python 3 on the Nano

I would also like that link. I have been trying to install CUDA all day on the nano hoping that it would help me install pytorch, but no luck.

Are you using the provided SD card image? CUDA should be easy to install that way. PyTorch takes a lot of memory and time to compile, though. I built it on Xavier for the other Jetsons.

I will upload the Python2 wheel and build directions later. If you want Python3 support you should be able to follow those steps to re-build. I will also try to get the Python3 wheel built and uploaded this week.

Hi Dustin,

I have tried the provided SD card image, and then tried to compile and install Pytorch master branch.

It encountered the following problems.

File "setup.py", line 710, in <module>
build_deps()
File "setup.py", line 282, in build_deps
build_dir = 'build')
File "/home/newton/pytorch/tools/build_pytorch_libs.py", line 273, in build_caffe2
check_call(['make', '-j', str(max_jobs), 'install'], cwd=build_dir, env=my_env)
FIle "/usr/lib/python3.6/subprocess.py", line 291, in check_call
raise CalledProcessError(retcode, cmd)
subprocess.CalledProcessError: Command '[' make', '-j', '4', 'install'] returned non-zero exist status 2.

Can you provide the link to the python 2.7 wheel for pytorch right now?

OK, here it is: https://nvidia.box.com/shared/static/m6vy0c7rs8t1alrt9dqf7yt1z587d1jk.whl

torch-1.1.0a0+b457266-cp27-cp27mu-linux_aarch64.whl

I tried it out. It gave me the following error:

torch-1.1.0a0+b457266-cp27-cp27mu-linux_aarch64.whl is not a supported wheel on this platform.

I tried to ssh connect to the board with ssh to avoid memory error:

pip --no-cache-dir install torch-1.1.0a0+b457266-cp27-cp27mu-linux_aarch64.whl

I tried to enlarge the swap memory file size, then i got this not supported error.

I have used the Jetpack SD official image:

4.9.140-tegra #1 SMP PREEMPT aarch64 aarch64 aarch64 GNU/Linux

I am not sure if yours are built with CUDA 10. Any suggestion?

Thanks Dustin.

Hi Dustin,

I think I could help with this too, could you let me know what changes you made on Pytorch cmake configuration to get this wheel? Also, what is the version of Jetpack you are using?

The latest one with the official SD image is Jet pack 4.2.

Can you run cat ~/.pip/pip.log?

That wheel was built for CUDA10, and haven’t had that problem installing it on other Nano’s using JetPack 4.2, but maybe it needs rebuilt for some reason? You shouldn’t need the swap just to install it, unless you are building from source.

It would seem you are using Python 2.7 with this - is anyone else able to install it? I’ll post the directions tomorrow.

Hi Dustin,

Thanks a lot! I confirm now it works for me on Nvidia Jetson Nano with JetPack 4.2 installed.

Previous error was due to the setting in python 3.6 pip. After switching from python 3.6 to python 2.7 completely and also install the future (pip2 install future), the Pytorch wheel you provided works for me.

I see, very good - that makes sense. Thanks for confirming that it works for you!

OK, I’ve built the PyTorch wheel for Python 3.6, here it is:

https://nvidia.box.com/shared/static/veo87trfaawj5pfwuqvhl6mzc5b55fbj.whl
torch-1.1.0a0+b457266-cp36-cp36m-linux_aarch64.whl

And I am preparing the build instructions for these to post.

Hi Dusty,

I just tried installing this with the following commands, and it seemed to work:

wget https://nvidia.box.com/shared/static/veo87trfaawj5pfwuqvhl6mzc5b55fbj.whl

mv https://nvidia.box.com/shared/static/veo87trfaawj5pfwuqvhl6mzc5b55fbj.whl torch-1.1.0a0+b457266-cp36-cp36m-linux_aarch64.whl

pip3 install torch-1.1.0a0+b457266-cp36-cp36m-linux_aarch64.whl

Thanks for turning this around so quickly!

I am not familiar with installing wheels, does this mean that pytorch is now installed if I was able to install the wheel?

To test that torch is installed OK, launch a python interpreter from terminal (python command for 2.7 or python3 command for 3.6), and run the following commands (ignore the <<< when typing in the commands):

>>> import pytorch
>>> print(torch.__version__)
>>> print('CUDA available: ' + str(torch.cuda.is_available()))
>>> a = torch.cuda.FloatTensor(2).zero_()
>>> print('Tensor a = ' + str(a))
>>> b = torch.randn(2).cuda()
>>> print('Tensor b = ' + str(b))
>>> c = a + b
>>> print('Tensor c = ' + str(c))

Hi all, I’ve provided the build instructions in this thread here: https://devtalk.nvidia.com/default/topic/1049071/jetson-nano/pytorch-for-jetson-nano/

@dusty_nv, Thank you very much.

I downloaded and installed your python3 wheel on my jetson-tx2. It looks good after verified the pytorch from the command in you answer above.

But when I ran my programme, it failed as follows:

File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/parallel/data_parallel.py", line 150, in forward
    return self.module(*inputs[0], **kwargs[0])
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/nvidia/test/ESPNetv2-sacmehta/segmentation/cnn/SegmentationModel.py", line 57, in forward
    up_l4_to_l3 = F.interpolate(out_l4_proj, scale_factor=2, mode='bilinear', align_corners=True)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 2507, in interpolate
    return torch._C._nn.upsample_bilinear2d(input, _output_size(2), align_corners)
RuntimeError: cuda runtime error (7) : too many resources requested for launch at /media/nvidia/WD_BLUE_2.5_1TB/pytorch/aten/src/THCUNN/generic/SpatialUpSamplingBilinear.cu:67

I have tried “sudo rm -rf ~/.nv && sudo reboot” but no luck.
I am re-building pytorch from source and would let you know the update.

My env:
jetson-tx2
jetpack4.2
cat /usr/local/lib/python3.6/dist-packages/torch/version.py
version = ‘1.1.0a0+b457266’
debug = False
cuda = ‘10.0.117’

nvidia@nvidia-tx2:~$ df -h
Filesystem Size Used Avail Use% Mounted on
/dev/root 28G 12G 15G 45% /
devtmpfs 3.8G 0 3.8G 0% /dev
tmpfs 3.9G 4.0K 3.9G 1% /dev/shm
tmpfs 3.9G 20M 3.9G 1% /run
tmpfs 5.0M 4.0K 5.0M 1% /run/lock
tmpfs 3.9G 0 3.9G 0% /sys/fs/cgroup
tmpfs 786M 12K 786M 1% /run/user/120
tmpfs 786M 108K 786M 1% /run/user/1000

Thanks.