PyTorch-安装 & 代码自动补全解决方案

最新推荐文章于 2025-05-19 16:50:36 发布

原创最新推荐文章于 2025-05-19 16:50:36 发布 · 1.8k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#PyTorch

PyTorch 专栏收录该内容

3 篇文章

订阅专栏

本文介绍如何在PyCharm中为PyTorch添加类型提示，实现代码自动补全功能，解决因PyTorch大部分代码由C++实现导致的类型推导问题。通过创建.pyi文件，可以不改变PyTorch源码的情况下，增强IDE的代码理解能力。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pip3 install torch torchvision  安装命令

# 构造一个未初始化的5x3矩阵
x = torch.empty(5, 3)
print(x)
# 运行结果
tensor([[ 0.0000e+00,  2.0000e+00,  1.3795e+15],
        [ 8.5920e+09, -1.8056e-34,  4.5745e-41],
        [-3.9593e-33,  4.5745e-41, -4.9454e-33],
        [ 4.5745e-41,  1.6954e-29,  1.6428e+27],
        [ 1.3790e+15, -2.5250e-29,  1.3789e+15]])

安装完PyTorch后，运行程序输出正常，但是发现PyCharm却有错误提示，what the fuck？
答案就是：由于PyTorch的绝大部分代码使用C++实现，因此其Python接口都是在import torch执行之后动态加载进来的，所以类型推导在执行之前是无法知道torch模块中具体有哪些函数、每个函数的参数类型和返回值了。

那么如何解决这个问题，实现自动补全函数？
方法就是写一个.pyi文件。这个文件里把你的函数、类的输入输出等定义好，PyCharm会自动识别。而在Python运行时又会直接忽略这个文件。因此我们就可以在完全不影响PyTorch本身的情况下为其增加类型提示。

所幸GitHub上已经有人写好了PyTorch类型提示的文件(需翻墙)
https://gist.github.com/t-vi/0d0ae013072f96f50fa11fbc2287e33b
文件存放路径：${PYTHON_PATH}/site-packages/torch/_init_.pyi