Pytorch 使用setup.py在CUDA上安装PyTorch

在本文中，我们将介绍如何使用setup.py在CUDA上安装PyTorch。PyTorch是一个基于Python的科学计算库，它广泛应用于机器学习和深度学习研究。CUDA是英伟达开发的一种通用并行计算架构，它可以提供强大的加速计算能力。通过在PyTorch中启用CUDA支持，我们可以利用计算机的GPU来加速深度学习任务的训练和推断过程。

阅读更多：Pytorch 教程

安装PyTorch

首先，我们需要在计算机上安装PyTorch。PyTorch提供了多种安装方式，包括使用pip、conda或者从源代码编译。其中，使用pip安装是最方便的方法之一。

pip install torch

在安装过程中，PyTorch将会自动判断系统上是否存在CUDA，并根据相应的CUDA版本编译相关的代码。如果计算机上没有安装CUDA，PyTorch会使用CPU进行计算。

检查CUDA支持

安装完成后，我们可以通过以下步骤来检查PyTorch是否成功启用了CUDA支持：

import torch

if torch.cuda.is_available():
    print("CUDA is available!")
    print("PyTorch is using CUDA version:", torch.version.cuda)
else:
    print("CUDA is not available!")

如果输出结果中显示“CUDA is available!”，表示PyTorch成功启用了CUDA支持，并且打印出当前使用的CUDA版本。

使用setup.py进行安装

如果我们希望在自己的项目中使用PyTorch，并想在setup.py中指定PyTorch的安装方式和前提条件，则可以通过以下步骤进行安装：

在项目的根目录下创建一个名为setup.py的文件。
在setup.py中导入所需的库和模块。

import setuptools

setuptools.setup(
    # ...
)

在setup()函数中设置安装参数和依赖项。

setuptools.setup(
    # ...
    install_requires=[
        'torch==1.7.0',
    ],
    extras_require={
        'cuda': ['torch==1.7.0+cu110']
    },
)

在上面的示例中，我们设置了对不同版本的PyTorch的依赖需求。默认情况下，安装时将使用torch==1.7.0，即不使用CUDA。如果想启用CUDA支持，可以使用extras_require来指定cuda选项，它将安装带有CUDA支持的PyTorch：torch==1.7.0+cu110。

运行 python setup.py install 来安装PyTorch。

示例

下面是一个完整的示例，展示了如何使用setup.py来安装与CUDA兼容的PyTorch：

import setuptools

setuptools.setup(
    name='my_package',
    version='0.1',
    author='Your Name',
    author_email='your@email.com',
    description='A sample package',
    packages=setuptools.find_packages(),
    install_requires=[
        'torch==1.7.0',
    ],
    extras_require={
        'cuda': ['torch==1.7.0+cu110']
    },
)