如何让我的 PyTorch DCGAN 代码在 GPU 上运行？

Question

我正在尝试在 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")

Answer 1

关于您的代码的一些问题：

1. device 变量的值是多少？
   确保是 torch.device('cuda:0')（或任何你的 GPU 设备 ID）。
   否则，如果您的 device 实际上是 torch.device('cpu')，那么您在 CPU 中运行。
   请参阅torch.device 了解更多信息。

2. 您删除了代码的“模型”部分，但您可能跳过了其中的一个重要部分：您是否也将模型移至 GPU？通常一个模型包含许多内部参数（也就是可训练的权重），您在设备上也需要它们。
   你的代码也应该有

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

请注意，与torch.tensors 不同，在nn.Module 上调用.to 是an "in place" operation。

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)

您还可以将torch.zeros_like 和torch.ones_like 用于target 变量：

target = torch.zeros_like(input)

请注意，zeros_like 和 one_like 会为您处理好设备（和数据类型） - 它与 input 的相同。