Python day34

赵英英俊2025-08-05 17:16

内容：

使用GPU训练及类的call方法
使用GPU训练：调用.to(device)就可以使用GPU进行训练，且只有继承nn.Module以及torch.Tensor类型的才可以调用上述方法，同时计算时所有的输入张量和模型必须处于用一个设备，否则会触发运行时错误
__call__方法：类内部定义的一个方法，使用实例化后的类名即可直接调用

代码：

python 复制代码

import torch

# 检查CUDA是否可用
if torch.cuda.is_available():
    print("CUDA可用！")
    # 获取可用的CUDA设备数量
    device_count = torch.cuda.device_count()
    print(f"可用的CUDA设备数量: {device_count}")
    # 获取当前使用的CUDA设备索引
    current_device = torch.cuda.current_device()
    print(f"当前使用的CUDA设备索引: {current_device}")
    # 获取当前CUDA设备的名称
    device_name = torch.cuda.get_device_name(current_device)
    print(f"当前CUDA设备的名称: {device_name}")
    # 获取CUDA版本
    cuda_version = torch.version.cuda
    print(f"CUDA版本: {cuda_version}")
    # 查看cuDNN版本（如果可用）
    print("cuDNN版本:", torch.backends.cudnn.version())

else:
    print("CUDA不可用。")


# 设置GPU设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")


# 实例化模型并移至GPU
model = MLP().to(device)

python 复制代码

class Studen():
    def __init__(self, name):
        self.name = name
    def __call(self):
        print("Hello + self.name")

stu1 = Student()
stu1() # 输出 Hello, name