C# 深度学习：对抗生成网络(GAN)训练头像生成模型

using dcgan;
using Maomi.Torch;
using System.Diagnostics;
using TorchSharp;
using TorchSharp.Modules;
using static TorchSharp.torch;

// 使用 GPU 启动
Device defaultDevice = MM.GetOpTimalDevice();
torch.set_default_device(defaultDevice);

// Set random seed for reproducibility
var manualSeed = 999;

// manualSeed = random.randint(1, 10000) # use if you want new results
Console.WriteLine("Random Seed:" + manualSeed);
random.manual_seed(manualSeed);
torch.manual_seed(manualSeed);

Options options = new Options()
{
    Dataroot = "E:\\datasets\\celeba",
    // 设置这个可以并发加载数据集，加快训练速度
    Workers = 10,
    BatchSize = 128,
};

public class Options
{
    /// <summary>
    /// Root directory for dataset
    /// </summary>
    public string Dataroot { get; set; } = "data/celeba";

    /// <summary>
    /// Number of workers for dataloader
    /// </summary>
    public int Workers { get; set; } = 2;

    /// <summary>
    /// Batch size during training
    /// </summary>
    public int BatchSize { get; set; } = 128;

    /// <summary>
    /// Spatial size of training images. All images will be resized to this size using a transformer.
    /// </summary>
    public int ImageSize { get; set; } = 64;

    /// <summary>
    /// Number of channels in the training images. For color images this is 3
    /// </summary>
    public int Nc { get; set; } = 3;

    /// <summary>
    /// Size of z latent vector (i.e. size of generator input)
    /// </summary>
    public int Nz { get; set; } = 100;

    /// <summary>
    /// Size of feature maps in generator
    /// </summary>
    public int Ngf { get; set; } = 64;

    /// <summary>
    /// Size of feature maps in discriminator
    /// </summary>
    public int Ndf { get; set; } = 64;

    /// <summary>
    /// Number of training epochs
    /// </summary>
    public int NumEpochs { get; set; } = 5;

    /// <summary>
    /// Learning rate for optimizers
    /// </summary>
    public double Lr { get; set; } = 0.0002;

    /// <summary>
    /// Beta1 hyperparameter for Adam optimizers
    /// </summary>
    public double Beta1 { get; set; } = 0.5;

    /// <summary>
    /// Number of GPUs available. Use 0 for CPU mode.
    /// </summary>
    public int Ngpu { get; set; } = 1;
}

/path/to/celeba
    -> img_align_celeba  
        -> 188242.jpg
        -> 173822.jpg
        -> 284702.jpg
        -> 537394.jpg
           ...

// 创建一个 samples 目录用于输出训练过程中产生的输出效果
if(Directory.Exists("samples"))
{
    Directory.Delete("samples", true);
}

Directory.CreateDirectory("samples");

// 加载图像并对图像做转换处理
var dataset = MM.Datasets.ImageFolder(options.Dataroot, torchvision.transforms.Compose(
    torchvision.transforms.Resize(options.ImageSize),
    torchvision.transforms.CenterCrop(options.ImageSize),
    torchvision.transforms.ConvertImageDtype(ScalarType.Float32),
    torchvision.transforms.Normalize(new double[] { 0.5, 0.5, 0.5 }, new double[] { 0.5, 0.5, 0.5 }))
);

// 分批加载图像
var dataloader = torch.utils.data.DataLoader(dataset, batchSize: options.BatchSize, shuffle: true, num_worker: options.Workers, device: defaultDevice);

var netG = new dcgan.Generator(options).to(defaultDevice);

static void weights_init(nn.Module m)
{
    var classname = m.GetType().Name;
    if (classname.Contains("Conv"))
    {
        if (m is Conv2d conv2d)
        {
            nn.init.normal_(conv2d.weight, 0.0, 0.02);
        }
    }
    else if (classname.Contains("BatchNorm"))
    {
        if (m is BatchNorm2d batchNorm2d)
        {
            nn.init.normal_(batchNorm2d.weight, 1.0, 0.02);
            nn.init.zeros_(batchNorm2d.bias);
        }
    }
}

public class Generator : nn.Module<Tensor, Tensor>, IDisposable
{
    private readonly Options _options;

    public Generator(Options options) : base(nameof(Generator))
    {
        _options = options;
        main = nn.Sequential(
            // input is Z, going into a convolution
            nn.ConvTranspose2d(options.Nz, options.Ngf * 8, 4, 1, 0, bias: false),
            nn.BatchNorm2d(options.Ngf * 8),
            nn.ReLU(true),
            // state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(options.Ngf * 8, options.Ngf * 4, 4, 2, 1, bias: false),
            nn.BatchNorm2d(options.Ngf * 4),
            nn.ReLU(true),
            // state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(options.Ngf * 4, options.Ngf * 2, 4, 2, 1, bias: false),
            nn.BatchNorm2d(options.Ngf * 2),
            nn.ReLU(true),
            // state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(options.Ngf * 2, options.Ngf, 4, 2, 1, bias: false),
            nn.BatchNorm2d(options.Ngf),
            nn.ReLU(true),
            // state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(options.Ngf, options.Nc, 4, 2, 1, bias: false),
            nn.Tanh()
            // state size. (nc) x 64 x 64
            );

        RegisterComponents();
    }

    public override Tensor forward(Tensor input)
    {
        return main.call(input);
    }

    Sequential main;
}

var netG = new dcgan.Generator(options).to(defaultDevice);
netG.apply(weights_init);
Console.WriteLine(netG);

public class Discriminator : nn.Module<Tensor, Tensor>, IDisposable
{
    private readonly Options _options;

    public Discriminator(Options options) : base(nameof(Discriminator))
    {
        _options = options;

        main = nn.Sequential(
            // input is (nc) x 64 x 64
            nn.Conv2d(options.Nc, options.Ndf, 4, 2, 1, bias: false),
            nn.LeakyReLU(0.2, inplace: true),
            // state size. (ndf) x 32 x 32
            nn.Conv2d(options.Ndf, options.Ndf * 2, 4, 2, 1, bias: false),
            nn.BatchNorm2d(options.Ndf * 2),
            nn.LeakyReLU(0.2, inplace: true),
            // state size. (ndf*2) x 16 x 16
            nn.Conv2d(options.Ndf * 2, options.Ndf * 4, 4, 2, 1, bias: false),
            nn.BatchNorm2d(options.Ndf * 4),
            nn.LeakyReLU(0.2, inplace: true),
            // state size. (ndf*4) x 8 x 8
            nn.Conv2d(options.Ndf * 4, options.Ndf * 8, 4, 2, 1, bias: false),
            nn.BatchNorm2d(options.Ndf * 8),
            nn.LeakyReLU(0.2, inplace: true),
            // state size. (ndf*8) x 4 x 4
            nn.Conv2d(options.Ndf * 8, 1, 4, 1, 0, bias: false),
            nn.Sigmoid()
            );

        RegisterComponents();
    }

    public override Tensor forward(Tensor input)
    {
        var output =  main.call(input);

        return output.view(-1, 1).squeeze(1);
    }

    Sequential main;
}

var netD = new dcgan.Discriminator(options).to(defaultDevice);
netD.apply(weights_init);
Console.WriteLine(netD);

var criterion = nn.BCELoss();
var fixed_noise = torch.randn(new long[] { options.BatchSize, options.Nz, 1, 1 }, device: defaultDevice);
var real_label = 1.0;
var fake_label = 0.0;
var optimizerD = torch.optim.Adam(netD.parameters(), lr: options.Lr, beta1: options.Beta1, beta2: 0.999);
var optimizerG = torch.optim.Adam(netG.parameters(), lr: options.Lr, beta1: options.Beta1, beta2: 0.999);

var img_list = new List<Tensor>();
var G_losses = new List<double>();
var D_losses = new List<double>();

Console.WriteLine("Starting Training Loop...");

Stopwatch stopwatch = new();
stopwatch.Start();
int i = 0;
// For each epoch
for (int epoch = 0; epoch < options.NumEpochs; epoch++)
{
    foreach (var item in dataloader)
    {
        var data = item[0];

        netD.zero_grad();
        // Format batch
        var real_cpu = data.to(defaultDevice);
        var b_size = real_cpu.size(0);
        var label = torch.full(new long[] { b_size }, real_label, dtype: ScalarType.Float32, device: defaultDevice);
        // Forward pass real batch through D
        var output = netD.forward(real_cpu);
        // Calculate loss on all-real batch
        var errD_real = criterion.call(output, label);
        // Calculate gradients for D in backward pass
        errD_real.backward();
        var D_x = output.mean().item<float>();

        // Train with all-fake batch
        // Generate batch of latent vectors
        var noise = torch.randn(new long[] { b_size, options.Nz, 1, 1 }, device: defaultDevice);
        // Generate fake image batch with G
        var fake = netG.call(noise);
        label.fill_(fake_label);
        // Classify all fake batch with D
        output = netD.call(fake.detach());
        // Calculate D's loss on the all-fake batch
        var errD_fake = criterion.call(output, label);
        // Calculate the gradients for this batch, accumulated (summed) with previous gradients
        errD_fake.backward();
        var D_G_z1 = output.mean().item<float>();
        // Compute error of D as sum over the fake and the real batches
        var errD = errD_real + errD_fake;
        // Update D
        optimizerD.step();

        ////////////////////////////
        // (2) Update G network: maximize log(D(G(z)))
        ////////////////////////////
        netG.zero_grad();
        label.fill_(real_label);  // fake labels are real for generator cost
        // Since we just updated D, perform another forward pass of all-fake batch through D
        output = netD.call(fake);
        // Calculate G's loss based on this output
        var errG = criterion.call(output, label);
        // Calculate gradients for G
        errG.backward();
        var D_G_z2 = output.mean().item<float>();
        // Update G
        optimizerG.step();

        // ex: [0/25][4/3166] Loss_D: 0.5676 Loss_G: 7.5972 D(x): 0.9131 D(G(z)): 0.3024 / 0.0007
        Console.WriteLine($"[{epoch}/{options.NumEpochs}][{i%dataloader.Count}/{dataloader.Count}] Loss_D: {errD.item<float>():F4} Loss_G: {errG.item<float>():F4} D(x): {D_x:F4} D(G(z)): {D_G_z1:F4} / {D_G_z2:F4}");

        // 每处理 100 批，输出一次图片效果
        if (i % 100 == 0)
        {
            real_cpu.SaveJpeg("samples/real_samples.jpg");
            fake = netG.call(fixed_noise);
            fake.detach().SaveJpeg("samples/fake_samples_epoch_{epoch:D3}.jpg");
        }

        i++;
    }


    netG.save("samples/netg_{epoch}.dat");
    netD.save("samples/netd_{epoch}.dat");
}

Console.WriteLine("Training finished.");
stopwatch.Stop();
Console.WriteLine("Training Time: {stopwatch.Elapsed}");

netG.save("samples/netg.dat");
netD.save("samples/netd.dat");

Training finished.
Training Time: 00:49:45.6976041

Device defaultDevice = MM.GetOpTimalDevice();
torch.set_default_device(defaultDevice);

// Set random seed for reproducibility
var manualSeed = 999;

// manualSeed = random.randint(1, 10000) # use if you want new results
Console.WriteLine("Random Seed:" + manualSeed);
random.manual_seed(manualSeed);
torch.manual_seed(manualSeed);


Options options = new Options()
{
    Dataroot = "E:\\datasets\\celeba",
    Workers = 10,
    BatchSize = 128,
};

var netG = new dcgan.Generator(options);
netG.to(defaultDevice);
netG.load("netg.dat");

// 生成随机噪声
var fixed_noise = torch.randn(64, options.Nz, 1, 1, device: defaultDevice);

// 生成图像
var fake_images = netG.call(fixed_noise);

fake_images.SaveJpeg("fake_images.jpg");