目前 pytorch 已经升级到了 1.7.0 ,但在 ubuntu 20.04 下有一个非常诡异的 bug。为此,我们只能自己编译。
下面命令可以下载 pytorch :
git clone https://github.com/pytorch/pytorch
git checkout v1.7.0
git submodule sync
git submodule update --init --recursive
然后直接编译即可:
python setup.py build
编译好的库位于./build/lib.linux-x86_64-3.6/
,有 caffe2 和 torch 两个库。为了防止和系统库冲突,可以放在/opt/python
目录,然后设置环境变量(若想一直有效,该设置命令需放到/etc/environment
或者~/.zshrc
等命令行配置文件里):
export PYTHONPATH=/opt/python:$PYTHONPATH
接下来运行:
python -c 'import torch; print(torch.__file__)'
如果显示为/opt/python/torch/__init__.py
表示安装成功。
如果把编译好的库同步到其它机器,如果 GPU 不一样,可能出现下面错误:
>>> import torch
>>> x = torch.Tensor(3, 5, 6)
>>> x.to(torch.device("cuda:0"))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/opt/python/torch/tensor.py", line 179, in __repr__
return torch._tensor_str._str(self)
File "/opt/python/torch/_tensor_str.py", line 372, in _str
return _str_intern(self)
File "/opt/python/torch/_tensor_str.py", line 352, in _str_intern
tensor_str = _tensor_str(self, indent)
File "/opt/python/torch/_tensor_str.py", line 241, in _tensor_str
formatter = _Formatter(get_summarized_data(self) if summarize else self)
File "/opt/python/torch/_tensor_str.py", line 89, in __init__
nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
RuntimeError: CUDA error: no kernel image is available for execution on the device
这个是因为编译版本对 CUDA 架构的支持问题,默认只支持 7.5 (我也不知道这是啥),可以在编译前设置架构:
export TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5"
然后重新编译安装即可。
Q. E. D.