如何在 GPU 上将我的 PyTorch DCGAN 代码设为 运行?

How do i make my PyTorch DCGAN code to run on a GPU?

我正在尝试在 GPU 上训练 DCGAN,但是当我开始使用 PyTorch 时,我尝试从文档中做一些事情并且它有效但我想确认它是否是正确的方法我已经完成了,因为我在 GPU 上查看了许多关于 运行 它的其他问题,但它们是以不同的方式完成的。

import os
import torch
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch.nn import (
    BatchNorm2d,
    BCELoss,
    Conv2d,
    ConvTranspose2d,
    LeakyReLU,
    Module,
    ReLU,
    Sequential,
    Sigmoid,
    Tanh,
)
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.utils import save_image
from tqdm import tqdm

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

dataset = DataLoader(
    CIFAR10(
        root="./Data",
        download=True,
        transform=transforms.Compose(
            [
                transforms.Resize(64),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
            ]
        ),
    ),
    batch_size=64,
    shuffle=True,
    num_workers=2,
)

try:
    os.mkdir("./Models")
    os.mkdir("./Results")

except FileExistsError:
    pass


class Gen(Module):
    def __init__(self):
        super(Gen, self).__init__()
        self.main = Sequential(
            ConvTranspose2d(100, 512, 4, 1, 0, bias=False),
            BatchNorm2d(512),
            ReLU(True),
            ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
            BatchNorm2d(256),
            ReLU(True),
            ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            BatchNorm2d(128),
            ReLU(True),
            ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            BatchNorm2d(64),
            ReLU(True),
            ConvTranspose2d(64, 3, 4, 2, 1, bias=False),
            Tanh(),
        )

    def forward(self, input):
        output = self.main(input)
        return output


class Dis(Module):
    def __init__(self):
        super(Dis, self).__init__()
        self.main = Sequential(
            Conv2d(3, 64, 4, 2, 1, bias=False),
            LeakyReLU(0.2, inplace=True),
            Conv2d(64, 128, 4, 2, 1, bias=False),
            BatchNorm2d(128),
            LeakyReLU(0.2, inplace=True),
            Conv2d(128, 256, 4, 2, 1, bias=False),
            BatchNorm2d(256),
            LeakyReLU(0.2, inplace=True),
            Conv2d(256, 512, 4, 2, 1, bias=False),
            BatchNorm2d(512),
            LeakyReLU(0.2, inplace=True),
            Conv2d(512, 1, 4, 1, 0, bias=False),
            Sigmoid(),
        )

    def forward(self, input):
        output = self.main(input)
        return output.view(-1)


def weights(obj):
    classname = obj.__class__.__name__
    if classname.find("Conv") != -1:
        obj.weight.data.normal_(0.0, 0.02)

    elif classname.find("BatchNorm") != -1:
        obj.weight.data.normal_(1.0, 0.02)
        obj.bias.data.fill_(0)


gen = Gen()
gen.apply(weights)
gen.cuda().cuda().to(device)

dis = Dis()
dis.apply(weights)
dis.cuda().to(device)

criterion = BCELoss()

optimizerDis = Adam(dis.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizerGen = Adam(gen.parameters(), lr=0.0002, betas=(0.5, 0.999))

for epoch in range(25):
    for batch, data in enumerate(tqdm(dataset, total=len(dataset)), 0):
        dis.zero_grad()

        input = Variable(data[0]).cuda().to(device)
        target = Variable(torch.ones(input.size()[0])).cuda().to(device)
        output = dis(input).cuda().to(device)

        realError = criterion(output, target)

        noise = Variable(torch.randn(input.size()[0], 100, 1, 1)).cuda().to(device)
        fake = gen(noise).cuda().to(device)
        target = Variable(torch.zeros(input.size()[0])).cuda().to(device)
        output = dis(fake.detach()).cuda().to(device)

        fakeError = criterion(output, target)

        errD = realError + fakeError
        errD.backward()
        optimizerDis.step()

        gen.zero_grad()

        target = Variable(torch.ones(input.size()[0])).cuda().to(device)
        output = dis(fake).cuda().to(device)

        errG = criterion(output, target)
        errG.backward()
        optimizerGen.step()

        print(f"  {epoch+1}/25 Dis Loss: {errD.data:.4f} Gen Loss: {errG.data:.4f}")

save_image(data[0], "./Results/Real.png", normalize=True)
save_image(gen(noise).data, f"./Results/Fake{epoch+1}.png", normalize=True)
torch.save(gen, f"./Models/model{epoch+1}.pth")

关于您的代码的几点评论:

  1. 您的 device 变量的值是多少?
    确保是 torch.device('cuda:0')(或者你的 GPU 设备 ID 是什么)。
    否则,如果您的 device 实际上是 torch.device('cpu') 那么您 运行 在 CPU.
    有关详细信息,请参阅 torch.device

  2. 您删除了代码的 "model" 部分,但您可能跳过了那里的一个重要部分:您是否也将模型移至 GPU?通常一个模型包含许多内部参数(又名可训练权重),您在设备上也需要它们。
    你的代码也应该有

dis.to(device)
criterion.to(device)  # if your loss function also has trainable parameters

请注意,与 torch.tensors 不同,在 nn.Module 上调用 .to

  1. 你的代码有冗余:你不必调用 both .cuda() AND .to().
    调用 .cuda() 是将事物转移到 GPU 的旧方法,但是现在 pytorch 引入了 .to() 以使编码更简单。

  2. 由于您的输入和模型都在 GPU 上,因此您不需要明确地将输出也移动到设备。因此,您可以将 output = dis(input).cuda().to(device) 替换为 output = dis(input).

  3. 无需明确使用 Variable。你可以替换

noise = Variable(torch.randn(input.size()[0], 100, 1, 1)).cuda().to(device)

noise = torch.randn(input.size()[0], 100, 1, 1), device=input.device)

您还可以对 target 变量使用 torch.zeros_like and torch.ones_like

target = torch.zeros_like(input)

请注意,zeros_likeone_like 会为您处理设备(和数据类型)- 它将与 input 的相同。