python - 如何让我的 PyTorch DCGAN 代码在 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")

标签： pythonpytorchgpugenerative-adversarial-network

解决方案

关于您的代码的一些评论：

你的device变量值是多少？
确保是torch.device('cuda:0')（或者你的 GPU 的设备 ID 是什么）。
否则，如果您device实际上是，torch.device('cpu')那么您将在 CPU 中运行。
有关torch.device更多信息，请参阅。
您删除了代码的“模型”部分，但您可能跳过了其中的一个重要部分：您是否也将模型移到了 GPU 上？通常一个模型包含许多内部参数（也就是可训练的权重），你在设备上也需要它们。
你的代码也应该有

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

请注意，与torch.tensors 不同，调用.toann.Module是“就地”操作。

您的代码中有冗余：您不必同时调用 .cuda()AND .to()。
调用.cuda()是将事物转移到 GPU 的旧方法，但现在引入了 pytorch.to()以简化编码。
由于您的输入和模型都在 GPU 上，因此您也不需要将输出显式移动到设备。因此，您可以替换output = dis(input).cuda().to(device)为 just output = dis(input)。
无需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_likeandtorch.ones_like用于target变量：

target = torch.zeros_like(input)

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

问题描述

解决方案

推荐阅读