PyTorch之Tensor.to(device)和model.to(device)详解

温故而知新，可以为师矣！

一、参考资料

pytorch中Tensor.to(device)和model.to(device)的区别及说明

Pytorch to（device）

Pytorch的to(device)用法

二、PyTorch中的GPU/CUDA相关操作

import torch

# CUDA是否可用
torch.cuda.is_available()

# 获取GPU的数量
torch.cuda.device_count()

# 获取GPU的名称
torch.cuda.get_device_name(0)

pytorch.cuda

CUDA 语义

import torch

# 检查系统中是否有可用的 GPU
if torch.cuda.is_available():
    # 获取可用的 GPU 设备数量
    num_devices = torch.cuda.device_count()
    print("可用 GPU 数量:", num_devices)

    # 遍历所有可用的 GPU 设备并打印详细信息
    for i in range(num_devices):
        device = torch.cuda.get_device_properties(i)
        print(f"\nGPU {i} 的详细信息:")
        print("名称:", device.name)
        print("计算能力:", f"{device.major}.{device.minor}")
        print("内存总量 (GB):", round(device.total_memory / (1024**3), 1))
else:
    print("没有可用的 GPU")

示例一：

可用 GPU 数量: 1

GPU 0 的详细信息:
名称: NVIDIA L20
计算能力: 8.9
内存总量 (GB): 44.5

示例二：

可用 GPU 数量: 1

GPU 0 的详细信息:
名称: DCU K100_AI
计算能力: 9.2
内存总量 (GB): 64.0

三、PyTorch之to(device)

在PyTorch中，可以使用to()方法将Tensor或模型移动到指定的设备上。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
model.to(device)

以上两行代码放在读取数据之前。

mytensor = my_tensor.to(device)

这行代码的意思是，将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去，之后的运算都在GPU上进行。

1. Tensor.to(device)

功能：将Tensor移动到指定的设备上。

可以使用以下代码将Tensor移动到GPU上：

import torch


x = torch.randn(3, 3)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
x = x.to(device)

在这个示例中，首先创建一个3x3的Tensor x，然后使用 torch.device() 函数创建一个设备对象device。如果GPU可用，则将device设置为’cuda:0’，否则将device设置为’cpu’。最后，使用x.to(device) 将Tensor x移动到指定的设备上。

下述实例中，将 inputs和labels数据移动到指定的设备上。

import torch


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

data_transform = {
    "train": transforms.Compose([transforms.Resize((224, 224)),
                                 transforms.ToTensor(),
                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.255])]),
    "test": transforms.Compose([transforms.Resize((224, 224)),
                                transforms.ToTensor(),
                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.255])])}

# 训练集
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=data_transform['train'])
trainloader = DataLoader(trainset, batch_size=16, shuffle=True)

for i, data in enumerate(trainloader):
    # 获取输入
    inputs, labels = data
    inputs = inputs.to(device)  
    labels = labels.to(device)

2. model.to(device)

功能：将模型移动到指定的设备上。

可以使用以下代码将模型移动到GPU上：

import torch
import torch.nn as nn


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

    
model = Net()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(device)

在这个示例中，我们首先定义了一个简单的神经网络模型Net，并使用 net.to(device) 将模型移动到指定的设备上。其中，device=torch.device("cpu")表示使用cpu，而device=torch.device("cuda")则表示使用GPU。

如果有多个GPU，使用以下方法：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
model = Model()
 
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model，device_ids=[0,1,2])
  
model.to(device)

将由GPU保存的模型加载到GPU上。

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)

将由GPU保存的模型加载到CPU上。将torch.load()函数中的map_location参数设置为torch.device('cpu')。

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
model.to(device)

将由CPU保存的模型加载到GPU上。将torch.load()函数中的map_location参数设置为torch.device('cuda')。

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)

3. PyTorch中指定GPU

1. 直接终端中设定：

   export CUDA_VISIBLE_DEVICES=0
   

2. python代码中设定：

   import os
   os.environ['CUDA_VISIBLE_DEVICES'] = '0'
   

3. 使用set_device()函数：

   import torch
   
   
   # Pytoch中的`in-place operation`原地操作符
   torch.cuda.set_device(id)
   

4. 相关经验

• Tensor 和 Numpy都是矩阵，区别是前者可以在GPU上运行，后者只能在CPU上；
• Tensor可以直接通过print显示数据类型，而Numpy不可以；

四、常用函数

total_memory

功能：计算总显存。

import torch

# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        props = torch.cuda.get_device_properties(i)
        g_gpu_memory = props.total_memory
        print(f"GPU{i}显存为: {g_gpu_memory / (1024 ** 3):.2f} GB")

memory_allocated()

功能：计算显存占用大小。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
gpu_memory_allocated = torch.cuda.memory_allocated(device)

print(f"\033[32m模型 qwen-vl-chat 加载完成，显存占用: {gpu_memory_allocated / 1024 ** 3:.2f} GB，剩余显存：{(g_gpu_memory - gpu_memory_allocated) / (1024 ** 3):.2f} GB\033[0m\n")

empty_cache()

功能：清理显存。

def load_and_run_model(image_url, prompt):
    # 释放显存
    history = []
    torch.cuda.empty_cache()
    print(f"\033[31m本模型较大，加载耗时约1分钟，请耐心等待...\033[0m")
    
    #加载模型
    try:
        print(f"开始加载 qwen-vl-chat 模型...")
        tokenizer = AutoTokenizer.from_pretrained("./qwen-vl-chat", trust_remote_code=True)

        # 打开bf16精度，A100、H100、RTX3060、RTX3070等显卡建议启用以节省显存
        # model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL-Chat", device_map="auto", trust_remote_code=True, bf16=True).eval()
        # 打开fp16精度，V100、P100、T4等显卡建议启用以节省显存
        # model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL-Chat", device_map="auto", trust_remote_code=True, fp16=True).eval()
        # 使用CPU进行推理，需要约32GB内存
        # model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL-Chat", device_map="cpu", trust_remote_code=True).eval()
        # 使用cuda进行推理
        model = AutoModelForCausalLM.from_pretrained("./qwen-vl-chat", device_map="cuda", trust_remote_code=True).eval()
        
        # Specify hyperparameters for generation
        model.generation_config = GenerationConfig.from_pretrained("./qwen-vl-chat", trust_remote_code=True)
        
        query = tokenizer.from_list_format([
            {'image': image_url}, # Either a local path or an url
            {'text': prompt},
        ])

        response, history = model.chat(tokenizer, query=query, history=None)
        print(response)
    except torch.cuda.OutOfMemoryError as e:
        print(f"\033[31m显存不足，无法继续运行模型...\033[0m")
        print(f"\033[31m错误信息: \033[0m")
        print(f"{str(e)}\n")
        # 释放显存
        torch.cuda.empty_cache()
        return False