[ pytorch ] 基本使用丨7. GPU分配丨

# 该代码转载自：https://www.cnblogs.com/darkknightzh/p/6836568.html
 
# 开始脚本，创建一个张量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
# 但是无论你获得一个新的Tensor或者Module
# 如果他们已经在目标设备上则不会执行复制操作
input = data.to(device)
model = MyModule(...).to(device)

这个方法需要注意的是，一定要记得在主程序里设置os.environ[“CUDA_VISIBLE_DEVICES”] ，不然cuda:0会找不到GPU

# ****************
#  主函数main.py
# ****************
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3, 4, 5"  # 将2, 3, 4, 5号GPU作为备选GPU
 
# 这样cuda:0表示使用 2 号GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

import torch
 
gpu_id = [0, 1, 2]
torch.cuda.set_device(id)
 
# 使用GPU
data.cuda()
model.cuda()

import torch.nn as nn
import torchvision
import os
import copy

class ft_net(nn.Module):
    def __init__(self, cfg=None):
        super(ft_net, self).__init__()
        ### 1. BackBone ###
        self.resnet50 = torchvision.models.resnet50(pretrained=True)

    def net(self, x):
        x =self.resnet50(x)
        return x

    def forward(self,x):
        x = self.resnet50(x)
        return x
os.environ['CUDA_VISIBLE_DEVICES'] = '1,0'

inputs = torch.rand([32,3, 224,224]).cuda()
model = ft_net().cuda()
model = torch.nn.DataParallel(model)

for i in range(10000):
    # output = model.module.net(inputs)   # 只能在一个GPU上运行
    output = model(inputs)                # 可以在两个GPU上运行