Collection of MindSpore implementations of Generative Adversarial Network varieties presented in research papers. 
"""AAE Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=10, help="dimensionality of the latent code")
parser.add_argument("--img_size",
                    type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


def reparameterization(mu, logvar):
    """reparameterization"""
    std = ops.exp(logvar / 2)
    sampled_z = ops.randn((mu.shape[0], opt.latent_dim))
    z = sampled_z * std + mu
    return z


class Encoder(nn.Cell):
    """Encoder Model"""

    def __init__(self):
        super().__init__(Encoder)

        self.model = nn.SequentialCell(
            nn.Dense(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2, ),
            nn.Dense(512, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2),
        )

        self.mu = nn.Dense(512, opt.latent_dim)
        self.logvar = nn.Dense(512, opt.latent_dim)

    def construct(self, img):
        img_flat = img.view(img.shape[0], -1)
        x = self.model(img_flat)
        mu = self.mu(x)
        logvar = self.logvar(x)
        z = reparameterization(mu, logvar)
        return z


class Decoder(nn.Cell):
    """Decoder Model"""

    def __init__(self):
        super().__init__(Decoder)

        self.model = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, int(np.prod(img_shape))),
            nn.Tanh(),
        )

    def construct(self, z):
        img_flat = self.model(z)
        img = img_flat.view(img_flat.shape[0], *img_shape)
        return img


class Discriminator(nn.Cell):
    """Discriminator Model"""

    def __init__(self):
        super().__init__(Discriminator)

        self.model = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dense(256, 1),
            nn.Sigmoid(),
        )

    def construct(self, z):
        validity = self.model(z)
        return validity


encoder = Encoder()
decoder = Decoder()
discriminator = Discriminator()
encoder.update_parameters_name("encoder")
decoder.update_parameters_name("decoder")
discriminator.update_parameters_name("discriminator")

# Use binary cross-entropy loss
adversarial_loss = nn.BCELoss()
pixelwise_loss = nn.L1Loss()

params = list(encoder.trainable_params()) + list(decoder.trainable_params())
optimizer_G = nn.optim.Adam(params, learning_rate=opt.lr,
                            beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr,
                            beta1=opt.b1, beta2=opt.b2)


def sample_image(n_row, batches):
    """Saves a grid of generated digits"""
    # Sample noise
    z = ops.randn((n_row ** 2, opt.latent_dim))
    gen_imgs = decoder(z)
    to_image(gen_imgs, os.path.join("images", F'{batches}.png'))


def g_forward(_real_imgs):
    """生成器正向传播"""
    _encoded_imgs = encoder(_real_imgs)
    decoded_imgs = decoder(_encoded_imgs)
    _g_loss = (0.001 * adversarial_loss(discriminator(_encoded_imgs), valid)
               + 0.999 * pixelwise_loss(decoded_imgs, real_imgs))
    return _g_loss, _encoded_imgs


def d_forward(_encoded_imgs, _valid, _fake):
    """判别器正向传播"""
    z = ops.randn((imgs.shape[0], opt.latent_dim))
    # Measure discriminator's ability to classify real from generated samples
    real_loss = adversarial_loss(discriminator(z), _valid)
    fake_loss = adversarial_loss(discriminator(_encoded_imgs), _fake)
    _d_loss = 0.5 * (real_loss + fake_loss)
    return _d_loss


transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# 训练
for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        encoder.set_train()
        decoder.set_train()
        discriminator.set_train()
        valid = ops.ones((imgs.shape[0], 1))
        fake = ops.zeros((imgs.shape[0], 1))
        ops.stop_gradient(valid)
        ops.stop_gradient(fake)

        real_imgs = Tensor(imgs)

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, encoded_imgs), g_grads = grad_g(real_imgs)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss), d_grads = grad_d(ops.stop_gradient(encoded_imgs), valid, fake)
        optimizer_D(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item()}] [G loss: {g_loss.asnumpy().item()}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            sample_image(10, batches_done)

implementations/aae/aae.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 10,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/aae/img_utils.py
mindspore-courses/MindSpore-GAN
"""ACGAN Model"""

import argparse
import os

import mindspore
import mindspore.common.initializer as init
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--n_classes",
                    type=int, default=10, help="number of classes for dataset")
parser.add_argument("--img_size",
                    type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.label_emb = nn.Embedding(opt.n_classes, opt.latent_dim)

        self.init_size = opt.img_size // 4  # Initial size before upsampling
        self.l1 = nn.SequentialCell(nn.Dense(opt.latent_dim, 128 * self.init_size ** 2))

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh(),
        )

    def construct(self, _noise, _labels):
        gen_input = ops.mul(self.label_emb(_labels), _noise)
        out = self.l1(gen_input)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            """Returns layers of each discriminator block"""
            block = [
                nn.Conv2d(in_filters, out_filters,
                          3, 2, 'pad', 1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0)))
            return block

        self.conv_blocks = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4

        # Output layers
        self.adv_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, 1),
            nn.Sigmoid()
        )
        self.aux_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, opt.n_classes),
            nn.Softmax())

    def construct(self, _img):
        out = self.conv_blocks(_img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        label = self.aux_layer(out)

        return validity, label


# Loss functions
adversarial_loss = nn.BCELoss()
auxiliary_loss = nn.CrossEntropyLoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()
generator.update_parameters_name('generator')
discriminator.update_parameters_name('discriminator')
generator.set_train()
discriminator.set_train()

optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

optimizer_G.update_parameters_name('optimizer_G')
optimizer_D.update_parameters_name('optimizer_D')


def sample_image(n_row, batches):
    """Saves a grid of generated digits"""
    # Sample noise
    _z = ops.randn((n_row ** 2, opt.latent_dim))
    # Get labels ranging from 0 to n_classes for n rows
    _labels = np.array([num for _ in range(n_row) for num in range(n_row)])
    _labels = Tensor(_labels)
    _gen_imgs = generator(_z, _labels)
    to_image(_gen_imgs, os.path.join("images", F'{batches}.png'))


def g_forward(_z, _gen_labels):
    """Generator forward function"""
    _gen_imgs = generator(_z, _gen_labels)
    # Loss measures generator's ability to fool the discriminator
    validity, pred_label = discriminator(_gen_imgs)
    _g_loss = 0.5 * (adversarial_loss(validity, valid) + auxiliary_loss(pred_label, _gen_labels))
    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _labels, _gen_labels, _valid, _fake):
    """Discriminator forward function"""
    # Loss for real images
    real_pred, _real_aux = discriminator(_real_imgs)
    d_real_loss = (adversarial_loss(real_pred, _valid) + auxiliary_loss(_real_aux, _labels)) / 2

    # Loss for fake images
    fake_pred, _fake_aux = discriminator(_gen_imgs)
    d_fake_loss = (adversarial_loss(fake_pred, _fake) + auxiliary_loss(_fake_aux, gen_labels)) / 2

    # Total discriminator loss
    _d_loss = (d_real_loss + d_fake_loss) / 2
    return _d_loss, _real_aux, _fake_aux


transform = [
    transforms.Rescale(1.0 / 255.0, 0),
    transforms.Resize(opt.img_size),
    transforms.Normalize([0.5], [0.5]),
    transforms.HWC2CHW()
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=True)

for epoch in range(opt.n_epochs):
    for i, (imgs, labels) in enumerate(dataset.create_tuple_iterator()):
        batch_size = imgs.shape[0]
        valid = ops.stop_gradient(ops.ones((batch_size, 1)))
        fake = ops.stop_gradient(ops.zeros((batch_size, 1)))

        real_imgs = Tensor(imgs, dtype=mstype.float32)
        labels = Tensor(labels, dtype=mstype.int32)

        z = ops.randn((batch_size, opt.latent_dim))
        gen_labels = ops.randint(0, opt.n_classes, (batch_size,), dtype=mstype.int32)

        # ---------------------
        #  Train Generator
        # ---------------------

        (g_loss, gen_imgs), g_grads = grad_g(z, gen_labels)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss, real_aux, fake_aux), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs),
                                                       labels, gen_labels, valid, fake)

        pred = np.concatenate([real_aux.asnumpy(), fake_aux.asnumpy()], axis=0)
        gt = np.concatenate([labels.asnumpy(), gen_labels.asnumpy()], axis=0)
        d_acc = np.mean(np.argmax(pred, axis=1) == gt)

        optimizer_D(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f},acc:{100 * d_acc:.2f}%] [G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            sample_image(10, batches_done)

implementations/acgan/acgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 10,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/acgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""BEGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
from mindspore import nn, Tensor
from mindspore import ops
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=62, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="number of image channels")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator network"""

    def __init__(self):
        super().__init__(Generator)

        self.init_size = opt.img_size // 4
        self.l1 = nn.SequentialCell(nn.Dense(opt.latent_dim, 128 * self.init_size ** 2))

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, _noise):
        out = self.l1(_noise)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        _img = self.conv_blocks(out)
        return _img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        # Upsampling
        self.down = nn.SequentialCell(
            nn.Conv2d(opt.channels, 64, 3,
                      2, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.ReLU()
        )
        # Fully-connected layers
        self.down_size = opt.img_size // 2
        down_dim = 64 * (opt.img_size // 2) ** 2
        self.fc = nn.SequentialCell(
            nn.Dense(down_dim, 32),
            nn.BatchNorm1d(32, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.ReLU(),
            nn.Dense(32, down_dim),
            nn.BatchNorm1d(down_dim,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.ReLU()
        )
        # Upsampling
        self.up = nn.SequentialCell(nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
                                    nn.Conv2d(64, opt.channels, 3, 1,
                                              'pad', 1, weight_init=init.Normal(0.02, 0.0)))

    def construct(self, img):
        out = self.down(img)
        out = self.fc(out.view(out.shape[0], -1))
        out = self.up(out.view(out.shape[0], 64, self.down_size, self.down_size))
        return out


# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)


def g_forward(_z):
    """Generator forward function"""
    # Generate a batch of images
    _gen_imgs = generator(_z)
    # Loss measures generator's ability to fool the discriminator
    _g_loss = ops.mean(ops.abs(discriminator(_gen_imgs) - _gen_imgs))

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    d_real = discriminator(_real_imgs)
    d_fake = discriminator(_gen_imgs)

    _d_loss_real = ops.mean(ops.abs(d_real - _real_imgs))
    _d_loss_fake = ops.mean(ops.abs(d_fake - _gen_imgs))
    _d_loss = _d_loss_real - k * _d_loss_fake

    return _d_loss, _d_loss_real, _d_loss_fake


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=True)

# ----------
#  Training
# ----------

generator.set_train()
discriminator.set_train()

# BEGAN hyper parameters
gamma = 0.75
lambda_k = 0.001
k = 0.0

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        # Configure input
        real_imgs = Tensor(imgs)

        # -----------------
        #  Train Generator
        # -----------------

        z = ops.randn((imgs.shape[0], opt.latent_dim))

        (g_loss, gen_imgs), g_grads = grad_g(z)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss, d_loss_real, d_loss_fake), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs))
        optimizer_D(d_grads)

        # ----------------
        # Update weights
        # ----------------

        diff = ops.mean(gamma * d_loss_real - d_loss_fake)

        # Update weight term for fake samples
        k = k + lambda_k * diff.asnumpy().item()
        k = min(max(k, 0), 1)  # Constraint to interval [0, 1]

        # Update convergence metric
        M = d_loss_real + ops.abs(diff)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}] '
            f'-- M: {M}, k: {k}'
        )

        batches_done = epoch * dataset.get_dataset_size() + i

        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/began/began.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/began/img_utils.py
mindspore-courses/MindSpore-GAN
"""BGAN Model"""
# Reference: https://wiseodd.github.io/techblog/2017/03/07/boundary-seeking-gan/


import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from mindspore import nn, Tensor
from mindspore import ops
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator Network"""
    def __init__(self):
        super().__init__(Generator)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, _z):
        _img = self.model(_z)
        _img = _img.view(_img.shape[0], *img_shape)
        return _img


class Discriminator(nn.Cell):
    """Discriminator Network"""
    def __init__(self):
        super().__init__(Discriminator)

        self.model = nn.SequentialCell(
            nn.Dense(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dense(256, 1),
            nn.Sigmoid()
        )

    def construct(self, _img):
        img_flat = _img.view(_img.shape[0], -1)
        _validity = self.model(img_flat)
        return _validity


def boundary_seeking_loss(y_pred):
    """
    Boundary seeking loss.
    Reference: https://wiseodd.github.io/techblog/2017/03/07/boundary-seeking-gan/
    """
    return 0.5 * ops.mean((ops.log(y_pred) - ops.log(1 - y_pred)) ** 2)


discriminator_loss = nn.BCELoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)


def g_forward(_z, _valid):
    """Generator forward function"""
    # Generate a batch of images
    _gen_imgs = generator(_z)
    # Loss measures generator's ability to fool the discriminator
    _g_loss = boundary_seeking_loss(discriminator(_gen_imgs))

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_loss = discriminator_loss(discriminator(_real_imgs), _valid)
    fake_loss = discriminator_loss(discriminator(_gen_imgs), _fake)
    _d_loss = (real_loss + fake_loss) / 2

    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# ----------
#  Training
# ----------

generator.set_train()
discriminator.set_train()

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs.shape[0], 1)))
        fake = ops.stop_gradient(ops.zeros((imgs.shape[0], 1)))

        # Configure input
        real_imgs = Tensor(imgs)

        # -----------------
        #  Train Generator
        # -----------------
        z = ops.randn((imgs.shape[0], opt.latent_dim))

        (g_loss, gen_imgs), g_grads = grad_g(z, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs), valid, fake)
        optimizer_D(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}]'
        )

        batches_done = epoch * dataset.get_dataset_size() + i

        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/bgan/bgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/bgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""BicycleGAN Model"""

import argparse
import datetime
import os
import sys
import time

import mindspore
import mindspore.common.dtype as mstype
from mindspore import nn
from mindspore import ops

from datasets import Edges2ShoesDataset
from img_utils import to_image
from models import Generator, MultiDiscriminator, Encoder

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="edges2shoes", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=8, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=128, help="size of image height")
parser.add_argument("--img_width", type=int, default=128, help="size of image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--latent_dim", type=int, default=8, help="number of latent codes")
parser.add_argument("--sample_interval", type=int, default=400, help="interval between saving generator samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
parser.add_argument("--lambda_pixel", type=float, default=10, help="pixelwise loss weight")
parser.add_argument("--lambda_latent", type=float, default=0.5, help="latent loss weight")
parser.add_argument("--lambda_kl", type=float, default=0.01, help="kullback-leibler loss weight")
opt = parser.parse_args()
print(opt)

os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

input_shape = (opt.channels, opt.img_height, opt.img_width)

# Loss functions
mae_loss = nn.L1Loss()

# Initialize generator, encoder and discriminators
generator = Generator(opt.latent_dim, input_shape)
encoder = Encoder(opt.latent_dim)
D_VAE = MultiDiscriminator(input_shape)
D_LR = MultiDiscriminator(input_shape)

generator.update_parameters_name("generator")
encoder.update_parameters_name("encoder")
D_VAE.update_parameters_name("D_VAE")
D_LR.update_parameters_name("D_LR")

if opt.epoch != 0:
    # Load pretrained models
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/generator_{opt.epoch}.ckpt', generator)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/encoder_{opt.epoch}.ckpt', encoder)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_VAE_{opt.epoch}.ckpt', D_VAE)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_LR_{opt.epoch}.ckpt', D_LR)

# Optimizers
optimizer_E = nn.optim.Adam(encoder.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D_VAE = nn.optim.Adam(D_VAE.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D_LR = nn.optim.Adam(D_LR.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

optimizer_E.update_parameters_name("optimizer_E")
optimizer_G.update_parameters_name("optimizer_G")
optimizer_D_VAE.update_parameters_name("optimizer_D_VAE")
optimizer_D_LR.update_parameters_name("optimizer_D_LR")

dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        input_shape=input_shape
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(opt.batch_size)

val_dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        input_shape=input_shape,
        mode='val'
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(8)


def sample_image(batches):
    """Saves a generated sample from the validation set"""
    generator.set_train(False)
    imgs = next(val_dataset.create_tuple_iterator())
    img_samples = None
    for img_A, _ in zip(imgs[0], imgs[1]):
        # Repeat input image by number of desired columns
        _real_A = img_A.view(1, *img_A.shape).tile((opt.latent_dim, 1, 1, 1))
        # Sample latent representations
        sampled_z = ops.randn(opt.latent_dim, opt.latent_dim, dtype=mstype.float32)
        # Generate samples
        fake_B = generator(_real_A, sampled_z)
        # Concatenate samples horisontally
        fake_B = ops.cat(list(fake_B), -1)
        img_sample = ops.cat((img_A, fake_B), -1)
        img_sample = img_sample.view(1, *img_sample.shape)
        # Concatenate with previous samples vertically
        img_samples = img_sample if img_samples is None else ops.cat((img_samples, img_sample), -2)
    to_image(img_samples, os.path.join(f'images/{opt.dataset_name}', f'{batches}.png'))
    generator.set_train()


def reparameterization(_mu, _logvar):
    """Reparameterization function"""
    std = ops.exp(_logvar / 2)
    sampled_z = ops.randn(_mu.shape[0], opt.latent_dim, dtype=mstype.float32)
    _z = sampled_z * std + _mu
    return _z


def ge_forward(_real_A, _real_B, _valid):
    """Encoder forward func"""
    # Produce output using encoding of B (cVAE-GAN)
    mu, logvar = encoder(_real_B)
    encoded_z = reparameterization(mu, logvar)
    fake_B = generator(_real_A, encoded_z)

    # Pixelwise loss of translated image by VAE
    _loss_pixel = mae_loss(fake_B, _real_B)
    # Kullback-Leibler divergence of encoded B
    _loss_kl = 0.5 * ops.sum(ops.exp(logvar) + mu ** 2 - logvar - 1)
    # Adversarial loss
    loss_VAE_GAN = D_VAE.compute_loss(fake_B, _valid)

    # ---------
    # cLR-GAN
    # ---------

    # Produce output using sampled z (cLR-GAN)
    sampled_z = ops.randn(_real_A.shape[0], opt.latent_dim, dtype=mstype.float32)
    _fake_B = generator(_real_A, sampled_z)
    # cLR Loss: Adversarial loss
    loss_LR_GAN = D_LR.compute_loss(_fake_B, _valid)

    # ----------------------------------
    # Total Loss (Generator + Encoder)
    # ----------------------------------

    _loss_GE = loss_VAE_GAN + loss_LR_GAN + opt.lambda_pixel * _loss_pixel + opt.lambda_kl * _loss_kl
    return _loss_GE, _loss_pixel, _loss_kl, fake_B, _fake_B


def g_forward(_real_A):
    """Generator warmup forward func"""
    # Produce output using sampled z (cLR-GAN)
    sampled_z = ops.randn(_real_A.shape[0], opt.latent_dim, dtype=mstype.float32)
    _fake_B = generator(_real_A, sampled_z)

    # Latent L1 loss
    _mu, _ = encoder(_fake_B)
    _loss_latent = opt.lambda_latent * mae_loss(_mu, sampled_z)

    return _loss_latent


def d_vae_forward(_real_B, _fake_B, _valid, _fake):
    """Discriminator forward function"""

    _loss_D_VAE = D_VAE.compute_loss(_real_B, _valid) + D_VAE.compute_loss(_fake_B, _fake)

    return _loss_D_VAE


def d_lr_forward(_real_B, _fake_B, _valid, _fake):
    """Discriminator forward function"""

    _loss_D_LR = D_LR.compute_loss(_real_B, _valid) + D_LR.compute_loss(_fake_B, _fake)

    return _loss_D_LR


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=False)
grad_d_vae = ops.value_and_grad(d_vae_forward, None, optimizer_D_VAE.parameters, has_aux=False)
grad_ge = ops.value_and_grad(ge_forward, None, optimizer_E.parameters, has_aux=True)
grad_d_lr = ops.value_and_grad(d_lr_forward, None, optimizer_D_LR.parameters, has_aux=False)

# ----------
#  Training
# ----------

generator.set_train()
encoder.set_train()
D_LR.set_train()
D_VAE.set_train()

# Adversarial loss
valid = 1
fake = 0

prev_time = time.time()
for epoch in range(opt.n_epochs):
    # Model inputs
    for i, (real_A, real_B) in enumerate(dataset.create_tuple_iterator()):

        # -------------------------------
        #  Train Generator and Encoder
        # -------------------------------

        (loss_GE, loss_pixel, loss_kl, fake_B1, fake_B2), ge_grads = grad_ge(real_A, real_B, valid)
        optimizer_E(ge_grads)

        (loss_latent), g_grads = grad_g(real_A)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (loss_D_VAE), d_vae_grads = grad_d_vae(real_B, ops.stop_gradient(fake_B1), valid, fake)
        optimizer_D_VAE(d_vae_grads)

        (loss_D_LR), d_lr_grads = grad_d_lr(real_B, ops.stop_gradient(fake_B2), valid, fake)
        optimizer_D_LR(d_lr_grads)

        # --------------
        #  Log Progress
        # --------------

        # Determine approximate time left
        batches_done = epoch * dataset.get_dataset_size() + i
        batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D VAE loss: {loss_D_VAE.asnumpy().item():.4f}, '
            f'LR loss:{loss_D_LR.asnumpy().item():.4f}] '
            f'[G loss: {loss_GE.asnumpy().item():.4f}, '
            f'pixel: {loss_pixel.asnumpy().item():.4f}, '
            f'kl: {loss_kl.asnumpy().item():.4f}, '
            f'latent:{loss_latent.asnumpy().item():.4f}] '
            f'ETA: {time_left}'
        )

        if batches_done % opt.sample_interval == 0:
            sample_image(batches_done)

        if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
            # Save model checkpoints
            mindspore.save_checkpoint(generator, f'saved_models/{opt.dataset_name}/generator_{epoch}.ckpt')
            mindspore.save_checkpoint(encoder, f'saved_models/{opt.dataset_name}/encoder_{epoch}.ckpt')
            mindspore.save_checkpoint(D_VAE, f'saved_models/{opt.dataset_name}/D_VAE_{epoch}.ckpt')
            mindspore.save_checkpoint(D_LR, f'saved_models/{opt.dataset_name}/D_LR_{epoch}.ckpt')

implementations/bicyclegan/bicyclegan.py
mindspore-courses/MindSpore-GAN
"""Dataset Preprocess"""

import glob
import os

import numpy as np
from PIL import Image
from mindspore.dataset.transforms import Compose
from mindspore.dataset.vision import transforms, Inter


class Edges2ShoesDataset():
    """Edges2Shoes Dataset"""

    def __init__(self, root, input_shape, mode="train"):
        self.transform = Compose(
            [
                transforms.Resize(input_shape[-2:], Inter.BICUBIC),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5], is_hwc=False)
            ]
        )

        self.files = sorted(glob.glob(os.path.join(root, mode) + "/*.*"))

    def __getitem__(self, index):
        img = Image.open(self.files[index % len(self.files)])
        w, h = img.size
        img_A = img.crop((0, 0, w / 2, h))
        img_B = img.crop((w / 2, 0, w, h))

        if np.random.random() < 0.5:
            img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], "RGB")
            img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], "RGB")

        img_A = self.transform(img_A)
        img_B = self.transform(img_B)

        return img_A[0], img_B[0]

    def __len__(self):
        return len(self.files)

implementations/bicyclegan/datasets.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/bicyclegan/img_utils.py
mindspore-courses/MindSpore-GAN
"""Networks"""
from mindcv import resnet18
from mindspore import nn, ops
import mindspore.common.initializer as init


##############################
#           U-NET
##############################


class UNetDown(nn.Cell):
    """U-Net Down"""

    def __init__(self, in_size, out_size, normalize=True):
        super().__init__(UNetDown)
        layers = [nn.Conv2d(in_size, out_size, 3, stride=2, pad_mode='pad', padding=1, has_bias=False,
                            weight_init=init.Normal(0.02, 0.0))]
        if normalize:
            layers.append(nn.BatchNorm2d(out_size, 0.8, affine=False,
                                         gamma_init=init.Normal(0.02, 1.0),
                                         beta_init=init.Constant(0.0)))
        layers.append(nn.LeakyReLU(0.2))
        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x)


class UNetUp(nn.Cell):
    """U-Net Up"""

    def __init__(self, in_size, out_size):
        super().__init__(UNetUp)
        self.model = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(in_size, out_size, 3, stride=1, pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(out_size, 0.8, affine=False,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.ReLU()
        )

    def construct(self, x, skip_input):
        x = self.model(x)
        x = ops.cat((x, skip_input), 1)
        return x


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self, latent_dim, img_shape):
        super().__init__(Generator)
        channels, self.h, self.w = img_shape

        self.fc = nn.Dense(latent_dim, self.h * self.w)

        self.down1 = UNetDown(channels + 1, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256)
        self.down4 = UNetDown(256, 512)
        self.down5 = UNetDown(512, 512)
        self.down6 = UNetDown(512, 512)
        self.down7 = UNetDown(512, 512, normalize=False)
        self.up1 = UNetUp(512, 512)
        self.up2 = UNetUp(1024, 512)
        self.up3 = UNetUp(1024, 512)
        self.up4 = UNetUp(1024, 256)
        self.up5 = UNetUp(512, 128)
        self.up6 = UNetUp(256, 64)

        self.final = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, channels, 3, stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, x, z):
        # Propogate noise through fc layer and reshape to img shape
        z = self.fc(z).view(z.shape[0], 1, self.h, self.w)
        d1 = self.down1(ops.cat((x, z), 1))
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        d7 = self.down7(d6)
        u1 = self.up1(d7, d6)
        u2 = self.up2(u1, d5)
        u3 = self.up3(u2, d4)
        u4 = self.up4(u3, d3)
        u5 = self.up5(u4, d2)
        u6 = self.up6(u5, d1)

        return self.final(u6)


##############################
#        Encoder
##############################


class Encoder(nn.Cell):
    """Encoder Network"""

    def __init__(self, latent_dim):
        super().__init__(Encoder)
        resnet18_model = resnet18(pretrained=False)
        self.feature_extractor = nn.SequentialCell(*list(resnet18_model.cells())[:-3])
        self.pooling = nn.AvgPool2d(kernel_size=8, stride=8, padding=0)
        # Output is mu and log(var) for reparameterization trick used in VAEs
        self.fc_mu = nn.Dense(256, latent_dim)
        self.fc_logvar = nn.Dense(256, latent_dim)

    def construct(self, img):
        out = self.feature_extractor(img)
        out = self.pooling(out)
        out = out.view(out.shape[0], -1)
        mu = self.fc_mu(out)
        logvar = self.fc_logvar(out)
        return mu, logvar


##############################
#        Discriminator
##############################


class MultiDiscriminator(nn.Cell):
    """Multi-Discriminator"""
    def __init__(self, input_shape):
        super().__init__(MultiDiscriminator)

        def discriminator_block(in_filters, out_filters, normalize=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.BatchNorm2d(out_filters, 0.8,
                                             gamma_init=init.Normal(0.02, 1.0),
                                             beta_init=init.Constant(0.0)))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        channels, _, _ = input_shape
        # Extracts discriminator models
        self.models = nn.CellList()
        for _ in range(3):
            self.models.append(
                nn.SequentialCell(
                    *discriminator_block(channels, 64, normalize=False),
                    *discriminator_block(64, 128),
                    *discriminator_block(128, 256),
                    *discriminator_block(256, 512),
                    nn.Conv2d(512, 1, 3, pad_mode='pad', padding=1,
                              weight_init=init.Normal(0.02, 0.0))
                )
            )

        self.downsample = nn.AvgPool2d(channels, stride=2, pad_mode='pad', padding=[1, 1], count_include_pad=False)

    def compute_loss(self, x, gt):
        """Computes the MSE between model output and scalar gt"""
        loss = sum([ops.mean((out - gt) ** 2) for out in self.construct(x)])
        return loss

    def construct(self, x):
        outputs = []
        for m in self.models:
            outputs.append(m(x))
            x = self.downsample(x)
        return outputs

implementations/bicyclegan/models.py
mindspore-courses/MindSpore-GAN
"""CCGAN Model"""

import argparse
import os

import mindspore
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.dataset import CelebADataset
from mindspore.dataset.vision import transforms, Inter

from img_utils import to_image
from models import Generator, Discriminator

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=8, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=128, help="size of each image dimension")
parser.add_argument("--mask_size", type=int, default=32, help="size of random mask")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=500, help="interval between image sampling")
opt = parser.parse_args()
print(opt)

input_shape = (opt.channels, opt.img_size, opt.img_size)

# Loss function
adversarial_loss = nn.MSELoss()

# Initialize generator and discriminator
generator = Generator(input_shape)
discriminator = Discriminator(input_shape)

def preprocess(_imgs):
    """Dataset preprocess func"""
    _imgs = transforms.Resize((opt.img_size, opt.img_size), Inter.BICUBIC)(_imgs)
    _imgs_lr = transforms.Resize((opt.img_size // 4, opt.img_size // 4), Inter.BICUBIC)(_imgs)
    _imgs = transforms.ToTensor()(_imgs)
    _imgs = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False)(_imgs)
    _imgs_lr = transforms.ToTensor()(_imgs_lr)
    _imgs_lr = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False)(_imgs_lr)

    return _imgs, _imgs_lr


dataset = CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True
).map(operations=[preprocess], input_columns=["image"], output_columns=["image", "image_lr"]).batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), opt.lr, opt.b1, opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), opt.lr, opt.b1, opt.b2)


def apply_random_mask(_imgs):
    """Add random masks to images."""
    idx = np.random.randint(0, opt.img_size - opt.mask_size, (_imgs.shape[0], 2))

    _masked_imgs = _imgs.copy().asnumpy()
    for j, (y1, x1) in enumerate(idx):
        y2, x2 = y1 + opt.mask_size, x1 + opt.mask_size
        _masked_imgs[j, :, y1:y2, x1:x2] = -1

    return Tensor(_masked_imgs)


def save_sample(_saved_samples):
    """Save samples."""
    # Generate inpainted image
    _gen_imgs = generator(_saved_samples["masked"], _saved_samples["lowres"])
    # Save sample
    sample = ops.cat((_saved_samples["masked"], _gen_imgs, _saved_samples["imgs"]), -2)
    to_image(sample, os.path.join(f'images/{opt.dataset_name}', F'{batches_done}.png'))


def g_forward(_masked_imgs, _imgs_lr, _valid):
    """Generator forward function"""
    # Generate a batch of images
    _gen_imgs = generator(_masked_imgs, _imgs_lr)

    # Loss measures generator's ability to fool the discriminator
    _g_loss = adversarial_loss(discriminator(_gen_imgs), _valid)
    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_loss = adversarial_loss(discriminator(_real_imgs), _valid)
    fake_loss = adversarial_loss(discriminator(_gen_imgs), _fake)
    _d_loss = 0.5 * (real_loss + fake_loss)
    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# ----------
#  Training
# ----------

saved_samples = {}

for epoch in range(opt.n_epochs):
    # Model inputs
    for i, batch in enumerate(dataset.create_dict_iterator()):
        imgs = batch["image"]
        imgs_lr = batch["image_lr"]

        masked_imgs = apply_random_mask(imgs)

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs.shape[0], *discriminator.output_shape)))
        fake = ops.stop_gradient(ops.zeros((imgs.shape[0], *discriminator.output_shape)))

        real_imgs = Tensor(imgs)
        imgs_lr = Tensor(imgs_lr)
        masked_imgs = Tensor(masked_imgs)

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(masked_imgs, imgs_lr, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (d_loss), d_grads = grad_d(real_imgs,
                                   ops.stop_gradient(gen_imgs), valid, fake)
        optimizer_D(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f} [G loss: {g_loss.asnumpy().item():.4f}]'
        )

        # Save first ten samples
        if not saved_samples:
            saved_samples["imgs"] = real_imgs[:1].copy()
            saved_samples["masked"] = masked_imgs[:1].copy()
            saved_samples["lowres"] = imgs_lr[:1].copy()
        elif saved_samples["imgs"].shape[0] < 10:
            saved_samples["imgs"] = ops.cat((saved_samples["imgs"], real_imgs[:1]), 0)
            saved_samples["masked"] = ops.cat((saved_samples["masked"], masked_imgs[:1]), 0)
            saved_samples["lowres"] = ops.cat((saved_samples["lowres"], imgs_lr[:1]), 0)

        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            save_sample(saved_samples)
        if batches_done % 5000 == 0:
            mindspore.save_checkpoint(generator,f'./g-{batches_done}.ckpt')
            mindspore.save_checkpoint(discriminator, f'./d-{batches_done}.ckpt')

implementations/ccgan/ccgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/ccgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""The Models of CCGAN"""

import mindspore.common.initializer as init
from mindspore import nn
from mindspore import ops


##############################
#           U-NET
##############################

class UNetDown(nn.Cell):
    """U-Net Down"""

    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
        super().__init__(UNetDown)
        layers = [
            nn.Conv2d(in_size, out_size, 4, 2,
                      'pad', 1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0))
        ]
        if normalize:
            layers.append(nn.BatchNorm2d(out_size, 0.8,
                                         gamma_init=init.Normal(0.02, 1.0),
                                         beta_init=init.Constant(0.0), affine=False))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout(p=dropout))

        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x)


class UNetUp(nn.Cell):
    """U-Net Up"""

    def __init__(self, in_size, out_size, dropout=0.0):
        super().__init__(UNetUp)
        layers = [
            nn.Conv2dTranspose(in_size, out_size, 4, 2,
                               'pad', 1, has_bias=False,
                               weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(out_size, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.ReLU()
        ]
        if dropout:
            layers.append(nn.Dropout(p=dropout))

        self.model = nn.SequentialCell(*layers)

    def construct(self, x, skip_input):
        x = self.model(x)
        x = ops.cat((x, skip_input), 1)

        return x


class Generator(nn.Cell):
    """U-Net Generator"""

    def __init__(self, input_shape):
        super().__init__(Generator)
        channels, _, _ = input_shape
        self.down1 = UNetDown(channels, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128 + channels, 256, dropout=0.5)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = UNetDown(512, 512, dropout=0.5)

        self.up1 = UNetUp(512, 512, dropout=0.5)
        self.up2 = UNetUp(1024, 512, dropout=0.5)
        self.up3 = UNetUp(1024, 256, dropout=0.5)
        self.up4 = UNetUp(512, 128)
        self.up5 = UNetUp(256 + channels, 64)

        final = [
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, channels, 3, 1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        ]
        self.final = nn.SequentialCell(*final)

    def construct(self, x, x_lr):
        # U-Net generator with skip connections from encoder to decoder
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d2 = ops.cat((d2, x_lr), 1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        u1 = self.up1(d6, d5)
        u2 = self.up2(u1, d4)
        u3 = self.up3(u2, d3)
        u4 = self.up4(u3, d2)
        u5 = self.up5(u4, d1)

        return self.final(u5)


##############################
#        Discriminator
##############################

class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, input_shape):
        super().__init__(Discriminator)

        channels, height, width = input_shape
        # Calculate output of image discriminator (PatchGAN)
        patch_h, patch_w = int(height / 2 ** 3), int(width / 2 ** 3)
        self.output_shape = (1, patch_h, patch_w)

        def discriminator_block(in_filters, out_filters, stride, normalize):
            """Returns layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 3, stride, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.InstanceNorm2d(out_filters,affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        layers = []
        in_filters = channels
        for out_filters, stride, normalize in [(64, 2, False), (128, 2, True), (256, 2, True), (512, 1, True)]:
            layers.extend(discriminator_block(in_filters, out_filters, stride, normalize))
            in_filters = out_filters

        layers.append(nn.Conv2d(out_filters, 1, 3, 1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)))

        self.model = nn.SequentialCell(*layers)

    def construct(self, img):
        return self.model(img)

implementations/ccgan/models.py
mindspore-courses/MindSpore-GAN
"""CGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--n_classes",
                    type=int, default=10, help="number of classes for dataset")
parser.add_argument("--img_size",
                    type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.label_emb = nn.Embedding(opt.n_classes, opt.n_classes)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim + opt.n_classes, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, noise, _labels):
        gen_input = ops.cat((self.label_emb(_labels), noise), -1)
        img = self.model(gen_input)
        img = img.view(img.shape[0], *img_shape)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        self.label_emb = nn.Embedding(opt.n_classes, opt.n_classes)

        self.model = nn.SequentialCell(
            nn.Dense(opt.n_classes + int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 512),
            nn.Dropout(0.4),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 512),
            nn.Dropout(0.4),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 1)
        )

    def construct(self, img, _labels):
        d_in = ops.cat((img.view(img.shape[0], -1),
                        self.label_emb(_labels)), -1)
        validity = self.model(d_in)
        return validity


# 损失函数
adversarial_loss = nn.MSELoss()

generator = Generator()
discriminator = Discriminator()

G_Optim = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
D_Optim = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_z, _gen_labels, _valid):
    """Generator forward function"""
    _gen_imgs = generator(_z, _gen_labels)
    validity = discriminator(_gen_imgs, _gen_labels)
    _g_loss = adversarial_loss(validity, _valid)
    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _labels, _gen_labels, _valid, _fake):
    """Discriminator forward function"""
    validity_real = discriminator(_real_imgs, _labels)
    d_real_loss = adversarial_loss(validity_real, _valid)

    validity_fake = discriminator(_gen_imgs, _gen_labels)
    d_fake_loss = adversarial_loss(validity_fake, _fake)

    _d_loss = (d_real_loss + d_fake_loss) / 2

    return _d_loss


transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, G_Optim.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, D_Optim.parameters, has_aux=False)

# 训练
for epoch in range(opt.n_epochs):
    generator.set_train()
    discriminator.set_train()
    for i, (imgs, labels) in enumerate(dataset.create_tuple_iterator()):
        batch_size = imgs.shape[0]

        valid = ops.ones((batch_size, 1))
        fake = ops.zeros((batch_size, 1))
        ops.stop_gradient(valid)
        ops.stop_gradient(fake)

        real_imgs = Tensor(imgs, dtype=mstype.float32)
        labels = Tensor(labels, dtype=mstype.int64)

        z = ops.randn((batch_size, opt.latent_dim))

        # gen_labels = Tensor(Tensor(np.random.randint(0, opt.n_classes, batch_size)), dtype=mstype.int64)
        gen_labels = ops.randint(0, opt.n_classes, (batch_size,))
        (g_loss, gen_imgs), g_grads = grad_g(z, gen_labels, valid)
        G_Optim(g_grads)

        (d_loss), d_grads = grad_d(real_imgs, gen_imgs, labels,
                                   gen_labels, valid, fake)
        D_Optim(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item()}] [G loss: {g_loss.asnumpy().item()}]'
        )

        batches_done = epoch * dataset.get_dataset_size() + i

        z = ops.randn((10 ** 2, opt.latent_dim))

        labels = np.array([num for _ in range(10) for num in range(10)])
        labels = Tensor(Tensor(labels), dtype=mstype.int64)
        gen_imgs = generator(z, labels)
        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs, os.path.join("images", F'{batches_done}.png'))

implementations/cgan/cgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 10,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/cgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""ClusterGAN Model"""

import argparse
import gzip
import itertools
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
import mindspore.numpy as mnp
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser(description="ClusterGAN Training Script")
parser.add_argument("-n", "--n_epochs", dest="n_epochs", default=200, type=int, help="Number of epochs")
parser.add_argument("-b", "--batch_size", dest="batch_size", default=64, type=int, help="Batch size")
parser.add_argument("-i", "--img_size", dest="img_size", type=int, default=28, help="Size of image dimension")
parser.add_argument("-d", "--latent_dim", dest="latent_dim", default=30, type=int, help="Dimension of latent space")
parser.add_argument("-l", "--lr", dest="learning_rate", type=float, default=0.0001, help="Learning rate")
parser.add_argument("-c", "--n_critic", dest="n_critic", type=int, default=5,
                    help="Number of training steps for discriminator per iter")
parser.add_argument("-w", "--wass_flag", dest="wass_flag", action='store_true', help="Flag for Wasserstein metric")
args = parser.parse_args()


def sample_z(shape=64, _latent_dim=10, _n_c=10, fix_class=-1):
    """Sample a random latent space vector"""
    assert (fix_class == -1 or (0 <= fix_class < _n_c)), f'Requested class {fix_class} outside bounds.'

    # Sample noise as generator input, zn
    zn = 0.75 * ops.stop_gradient(ops.randn((shape, _latent_dim), dtype=mstype.float32))

    ######### zc, zc_idx variables with grads, and zc to one-hot vector
    # Pure one-hot vector generation
    zc_FT = ops.zeros((shape, _n_c), dtype=mstype.float32)
    zc_idx = Tensor(mnp.empty(shape, dtype=mstype.int64))

    if fix_class == -1:
        zc_idx = ops.uniform(zc_idx.shape, minval=Tensor(0, dtype=mstype.int32),
                             maxval=Tensor(_n_c, dtype=mstype.int32), dtype=mstype.int32)
        zc_FT = zc_FT.scatter(1, zc_idx.unsqueeze(1), ops.ones((zc_idx.shape[0], 1), dtype=mstype.float32))
    else:
        zc_idx[:] = fix_class
        zc_FT[:, fix_class] = 1

        zc_FT = ops.stop_gradient(zc_FT)

    zc = zc_FT

    # Return components of latent space variable
    return zn, zc, zc_idx


def calc_gradient_penalty(D, real_data, generated_data):
    """Calculate gradient penalty"""
    # GP strength
    LAMBDA = 10

    b_size = real_data.shape[0]

    # Calculate interpolation
    alpha = ops.rand(b_size, 1, 1, 1)
    alpha = alpha.expand_as(real_data)

    interpolated = alpha * real_data + (1 - alpha) * generated_data

    # Calculate gradients of probabilities with respect to examples
    grad_fn = ops.grad(D)
    gradients = grad_fn(interpolated)

    # Gradients have shape (batch_size, num_channels, img_width, img_height),
    # so flatten to easily take norm per example in batch
    gradients = gradients.view(b_size, -1)

    # Derivatives of the gradient close to 0 can cause problems because of
    # the square root, so manually calculate norm and add epsilon
    gradients_norm = ops.sqrt(ops.sum(gradients ** 2, dim=1) + 1e-12)

    # Return gradient penalty
    return LAMBDA * ((gradients_norm - 1) ** 2).mean()


def softmax(x):
    """Softmax function"""
    return ops.softmax(x, axis=1)


class Reshape(nn.Cell):
    """
    Class for performing a reshape as a layer in a sequential model.
    """

    def __init__(self, shape=None):
        super().__init__(Reshape)
        if shape is None:
            shape = []
        self.shape = shape

    def construct(self, x):
        return x.view(x.shape[0], *self.shape)


class GeneratorCNN(nn.Cell):
    """
    CNN to model the generator of a ClusterGAN
    Input is a vector from representation space of dimension z_dim
    output is a vector from image space of dimension X_dim
    """

    # Architecture : FC1024_BR-FC7x7x128_BR-(64)4dc2s_BR-(1)4dc2s_S
    def __init__(self, _latent_dim, _n_c, _x_shape, verbose=False):
        super().__init__(GeneratorCNN)

        self.name = 'generator'
        self.latent_dim = _latent_dim
        self.n_c = _n_c
        self.x_shape = _x_shape
        self.ishape = (128, 7, 7)
        self.iels = int(np.prod(self.ishape))
        self.verbose = verbose

        self.model = nn.SequentialCell(
            # Fully connected layers
            nn.Dense(self.latent_dim + self.n_c, 1024,
                     weight_init=init.Normal(0.02, 0),
                     bias_init=init.Constant(0)),
            nn.BatchNorm1d(1024),
            nn.LeakyReLU(0.2),
            nn.Dense(1024, self.iels,
                     weight_init=init.Normal(0.02, 0),
                     bias_init=init.Constant(0)),
            nn.BatchNorm1d(self.iels),
            nn.LeakyReLU(0.2),

            # Reshape to 128 x (7x7)
            Reshape(self.ishape),

            # Upconvolution layers
            nn.Conv2dTranspose(128, 64, 4, stride=2,
                               pad_mode='pad', padding=1, has_bias=True,
                               weight_init=init.Normal(0.02, 0),
                               bias_init=init.Constant(0)),
            nn.BatchNorm2d(64, affine=False),
            nn.LeakyReLU(0.2),

            nn.Conv2dTranspose(64, 1, 4, stride=2,
                               pad_mode='pad', padding=1, has_bias=True,
                               weight_init=init.Normal(0.02, 0),
                               bias_init=init.Constant(0)),
            nn.Sigmoid()
        )

        if self.verbose:
            print(f'Setting up {self.name}...\n')
            print(self.model)

    def construct(self, zn, zc):
        z = ops.cat((zn, zc), 1)
        x_gen = self.model(z)
        # Reshape for output
        x_gen = x_gen.view(x_gen.shape[0], *self.x_shape)
        return x_gen


class EncoderCNN(nn.Cell):
    """
    CNN to model the encoder of a ClusterGAN
    Input is vector X from image space if dimension X_dim
    Output is vector z from representation space of dimension z_dim
    """

    def __init__(self, _latent_dim, _n_c, verbose=False):
        super().__init__(EncoderCNN)

        self.name = 'encoder'
        self.channels = 1
        self.latent_dim = _latent_dim
        self.n_c = _n_c
        self.cshape = (128, 5, 5)
        self.iels = int(np.prod(self.cshape))
        self.lshape = (self.iels,)
        self.verbose = verbose

        self.model = nn.SequentialCell(
            # Convolutional layers
            nn.Conv2d(self.channels, 64, 4, stride=2,
                      has_bias=True, pad_mode='valid',
                      weight_init=init.Normal(0.02, 0),
                      bias_init=init.Constant(0)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, 4, stride=2,
                      has_bias=True, pad_mode='valid',
                      weight_init=init.Normal(0.02, 0),
                      bias_init=init.Constant(0)),
            nn.LeakyReLU(0.2),

            # Flatten
            Reshape(self.lshape),

            # Fully connected layers
            nn.Dense(self.iels, 1024,
                     weight_init=init.Normal(0.02, 0),
                     bias_init=init.Constant(0)),
            nn.LeakyReLU(0.2),
            nn.Dense(1024, _latent_dim + _n_c,
                     weight_init=init.Normal(0.02, 0),
                     bias_init=init.Constant(0))
        )

        if self.verbose:
            print(f'Setting up {self.name}...\n')
            print(self.model)

    def construct(self, in_feat):
        z_img = self.model(in_feat)
        # Reshape for output
        z = z_img.view(z_img.shape[0], -1)
        # Separate continuous and one-hot components
        zn = z[:, 0:self.latent_dim]
        zc_logits = z[:, self.latent_dim:]
        # Softmax on zc component
        zc = softmax(zc_logits)
        return zn, zc, zc_logits


class DiscriminatorCNN(nn.Cell):
    """
    CNN to model the discriminator of a ClusterGAN
    Input is tuple (X,z) of an image vector and its corresponding
    representation z vector. For example, if X comes from the dataset, corresponding
    z is Encoder(X), and if z is sampled from representation space, X is Generator(z)
    Output is a 1-dimensional value
    """

    # Architecture : (64)4c2s-(128)4c2s_BL-FC1024_BL-FC1_S
    def __init__(self, _wass_metric=False, verbose=False):
        super().__init__(DiscriminatorCNN)

        self.name = 'discriminator'
        self.channels = 1
        self.cshape = (128, 5, 5)
        self.iels = int(np.prod(self.cshape))
        self.lshape = (self.iels,)
        self.wass = _wass_metric
        self.verbose = verbose

        self.model = nn.SequentialCell(
            # Convolutional layers
            nn.Conv2d(self.channels, 64, 4,
                      stride=2, has_bias=True, pad_mode='pad',
                      weight_init=init.Normal(0.02, 0),
                      bias_init=init.Constant(0)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, 4,
                      stride=2, has_bias=True, pad_mode='pad',
                      weight_init=init.Normal(0.02, 0),
                      bias_init=init.Constant(0)),
            nn.LeakyReLU(0.2),

            # Flatten
            Reshape(self.lshape),

            # Fully connected layers
            nn.Dense(self.iels, 1024,
                     weight_init=init.Normal(0.02, 0),
                     bias_init=init.Constant(0)),
            nn.LeakyReLU(0.2),
            nn.Dense(1024, 1,
                     weight_init=init.Normal(0.02, 0),
                     bias_init=init.Constant(0)),
        )

        # If NOT using Wasserstein metric, final Sigmoid
        if not self.wass:
            self.model = nn.SequentialCell(self.model, nn.Sigmoid())

        if self.verbose:
            print(f'Setting up {self.name}...\n')
            print(self.model)

    def construct(self, img):
        # Get output
        validity = self.model(img)
        return validity


# Training details
n_epochs = args.n_epochs
batch_size = args.batch_size
test_batch_size = 5000
lr = args.learning_rate
b1 = 0.5
b2 = 0.9
decay = 2.5 * 1e-5
n_skip_iter = args.n_critic

# Data dimensions
img_size = args.img_size
channels = 1

# Latent space info
latent_dim = args.latent_dim
n_c = 10
betan = 10
betac = 10

# Wasserstein+GP metric flag
wass_metric = args.wass_flag

x_shape = (channels, img_size, img_size)

# Loss function
bce_loss = nn.BCELoss()
xe_loss = nn.CrossEntropyLoss()
mse_loss = nn.MSELoss()

# Initialize generator and discriminator
generator = GeneratorCNN(latent_dim, n_c, x_shape)
encoder = EncoderCNN(latent_dim, n_c)
discriminator = DiscriminatorCNN(_wass_metric=wass_metric)

generator.update_parameters_name("generator")
encoder.update_parameters_name("encoder")
discriminator.update_parameters_name("discriminator")

param_G = list(generator.trainable_params()) + list(encoder.trainable_params())

ge_chain = itertools.chain(generator.trainable_params(),
                           encoder.trainable_params())

optimizer_GE = nn.optim.Adam(ge_chain, learning_rate=lr, beta1=b1, beta2=b2, weight_decay=decay)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=lr, beta1=b1, beta2=b2)

transform = [
    transforms.ToTensor(),
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir="../../data/MNIST",
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(batch_size, drop_remainder=True)

testdata = mindspore.dataset.MnistDataset(
    dataset_dir="../../data/MNIST",
    usage='test',
    shuffle=True
).map(operations=transform, input_columns="image").batch(batch_size)

test_imgs, test_labels = next(testdata.create_tuple_iterator())


def ge_forward(_real_imgs):
    """Generator forward function"""
    # Sample random latent variables
    zn, zc, zc_idx = sample_z(shape=_real_imgs.shape[0], _latent_dim=latent_dim, _n_c=n_c)
    # Generate a batch of images
    _gen_imgs = generator(zn, zc)
    # Discriminator output from real and generated samples
    D_gen = discriminator(_gen_imgs)
    # Encode the generated images
    enc_gen_zn, _, enc_gen_zc_logits = encoder(_gen_imgs)

    # Calculate losses for z_n, z_c
    zn_loss = mse_loss(enc_gen_zn, zn)
    zc_loss = xe_loss(enc_gen_zc_logits, zc_idx)

    # Check requested metric
    if wass_metric:
        # Wasserstein GAN loss
        _ge_loss = ops.mean(D_gen) + betan * zn_loss + betac * zc_loss
    else:
        # Vanilla GAN loss
        _valid = ops.stop_gradient(ops.ones((_gen_imgs.shape[0], 1), dtype=mstype.float32))
        v_loss = bce_loss(D_gen, _valid)
        _ge_loss = v_loss + betan * zn_loss + betac * zc_loss

    return _ge_loss, _gen_imgs, _valid


def d_forward(_real_imgs, _gen_imgs, _valid):
    """Discriminator forward function"""
    # Discriminator output from real and generated samples
    _D_gen = discriminator(_gen_imgs)
    _D_real = discriminator(real_imgs)
    # Measure discriminator's ability to classify real from generated samples
    if wass_metric:
        # Gradient penalty term
        grad_penalty = calc_gradient_penalty(discriminator, _real_imgs, _gen_imgs)

        # Wasserstein GAN loss w/gradient penalty
        _d_loss = ops.mean(_D_real) - ops.mean(_D_gen) + grad_penalty

    else:
        # Vanilla GAN loss
        fake = ops.stop_gradient(ops.zeros((_gen_imgs.shape[0], 1), dtype=mstype.float32))
        real_loss = bce_loss(_D_real, _valid)
        fake_loss = bce_loss(_D_gen, fake)
        _d_loss = (real_loss + fake_loss) / 2

    return _d_loss


grad_ge = ops.value_and_grad(ge_forward, None, optimizer_GE.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# ----------
#  Training
# ----------
ge_l = []
d_l = []

c_zn = []
c_zc = []
c_i = []

# Training loop
print(f'\nBegin training session with {n_epochs} epochs...\n')
for epoch in range(n_epochs):
    itruth_label = None
    for i, (imgs, label) in enumerate(dataset.create_tuple_iterator()):
        itruth_label = label
        # Ensure generator/encoder are trainable
        generator.set_train()
        discriminator.set_train()

        # Configure input
        real_imgs = imgs

        # ---------------------------
        #  Train Generator + Encoder
        # ---------------------------

        # Step for Generator & Encoder, n_skip_iter times less than for discriminator

        (ge_loss, gen_imgs, valid), ge_grads = grad_ge(real_imgs)
        if i % n_skip_iter == 0:
            optimizer_GE(ge_grads)
        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss), d_grads = grad_d(real_imgs, gen_imgs, valid)
        optimizer_D(d_grads)

    # Save training losses
    d_l.append(d_loss.asnumpy().item())
    ge_l.append(ge_loss.asnumpy().item())

    generator.set_train(False)
    discriminator.set_train(False)

    # Set number of examples for cycle calcs
    n_sqrt_samp = 5
    n_samp = n_sqrt_samp * n_sqrt_samp

    ## Cycle through test real -> enc -> gen
    t_imgs, t_label = test_imgs, test_labels
    # Encode sample real instances
    e_tzn, e_tzc, e_tzc_logits = encoder(t_imgs)
    # Generate sample instances from encoding
    teg_imgs = generator(e_tzn, e_tzc)
    # Calculate cycle reconstruction loss
    img_mse_loss = mse_loss(t_imgs, teg_imgs)
    # Save img reco cycle loss
    c_i.append(img_mse_loss.asnumpy().item())

    ## Cycle through randomly sampled encoding -> generator -> encoder
    zn_samp, zc_samp, zc_samp_idx = sample_z(shape=n_samp, _latent_dim=latent_dim, _n_c=n_c)
    # Generate sample instances
    gen_imgs_samp = generator(zn_samp, zc_samp)

    # Encode sample instances
    zn_e, zc_e, zc_e_logits = encoder(gen_imgs_samp)

    # Calculate cycle latent losses
    lat_mse_loss = mse_loss(zn_e, zn_samp)
    lat_xe_loss = xe_loss(zc_e_logits, zc_samp_idx)

    # Save latent space cycle losses
    c_zn.append(lat_mse_loss.asnumpy().item())
    c_zc.append(lat_xe_loss.asnumpy().item())

    # Save cycled and generated examples!
    r_imgs, i_label = real_imgs[:n_samp], itruth_label[:n_samp]
    e_zn, e_zc, e_zc_logits = encoder(r_imgs)
    reg_imgs = generator(e_zn, e_zc)

    to_image(reg_imgs, f'images/cycle_reg_{epoch:06}.png')
    to_image(gen_imgs_samp, f'images/gen_{epoch:06}.png')

    ## Generate samples for specified classes
    stack_imgs = []
    for idx in range(n_c):
        # Sample specific class
        zn_samp, zc_samp, zc_samp_idx = sample_z(shape=n_c, _latent_dim=latent_dim, _n_c=n_c, fix_class=idx)

        # Generate sample instances
        gen_imgs_samp = generator(zn_samp, zc_samp)

        if len(stack_imgs) == 0:
            stack_imgs = gen_imgs_samp
        else:
            stack_imgs = ops.cat((stack_imgs, gen_imgs_samp), 0)

    # Save class-specified generated examples!
    to_image(stack_imgs, f'images/gen_classes_{epoch:06}.png')

    print(f'[Epoch {epoch}/{n_epochs}] '
          f'Model Losses: [D: {d_loss.asnumpy().item():.4f}] '
          f'[GE: {ge_loss.asnumpy().item():.4f}]'
          )

    print(f'Cycle Losses: [x: {img_mse_loss.asnumpy().item():.4f}] '
          f'[z_n: {lat_mse_loss.asnumpy().item():.4f}] '
          f'[z_c: {lat_xe_loss.asnumpy().item():.4f}]'
          )

implementations/cluster_gan/cluster_gan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/cluster_gan/img_utils.py
mindspore-courses/MindSpore-GAN
"""COGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
from mindspore import nn
from mindspore import ops
from mindspore.dataset.transforms import Compose
from mindspore.dataset.vision import transforms

from img_utils import to_image
from mnistm import MNISTM

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=32, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class CoupledGenerators(nn.Cell):
    """COGAN Generators"""

    def __init__(self):
        super().__init__(CoupledGenerators)

        self.init_size = opt.img_size // 4
        self.fc = nn.SequentialCell(nn.Dense(opt.latent_dim, 128 * self.init_size ** 2))

        self.shared_conv = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
        )
        self.G1 = nn.SequentialCell(
            nn.Conv2d(128, 64, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )
        self.G2 = nn.SequentialCell(
            nn.Conv2d(128, 64, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, noise):
        out = self.fc(noise)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img_emb = self.shared_conv(out)
        img1 = self.G1(img_emb)
        img2 = self.G2(img_emb)
        return img1, img2


class CoupledDiscriminators(nn.Cell):
    """COGAN Discriminator"""

    def __init__(self):
        super().__init__(CoupledDiscriminators)

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [nn.Conv2d(in_filters, out_filters, 3, 2,
                               pad_mode='pad', padding=1, has_bias=False,
                               weight_init=init.Normal(0.02, 0.0))]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0), affine=False))
            block.extend([nn.LeakyReLU(0.2), nn.Dropout2d(0.25)])
            return block

        self.shared_conv = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )
        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4
        self.D1 = nn.Dense(128 * ds_size ** 2, 1)
        self.D2 = nn.Dense(128 * ds_size ** 2, 1)

    def construct(self, img1, img2):
        # Determine validity of first image
        out = self.shared_conv(img1)
        out = out.view(out.shape[0], -1)
        validity1 = self.D1(out)
        # Determine validity of second image
        out = self.shared_conv(img2)
        out = out.view(out.shape[0], -1)
        validity2 = self.D2(out)

        return validity1, validity2


# Loss function
adversarial_loss = nn.MSELoss()

# Initialize models
coupled_generators = CoupledGenerators()
coupled_discriminators = CoupledDiscriminators()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5],is_hwc=False)
]


os.makedirs("../../data/MNIST-M", exist_ok=True)

dataset1 = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

dataset2 = mindspore.dataset.GeneratorDataset(
    source=MNISTM(
        root='../../data/MNIST-M',
        mnist_root='../../data/MNIST',
        transform=Compose(
            [
                transforms.Resize(opt.img_size),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
            ]
        )
    ),
    shuffle=True,
    column_names=["image", "target"]
).batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(coupled_generators.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(coupled_discriminators.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_batch_size, _valid):
    """Generator forward function"""
    # Sample noise as generator input
    z = ops.randn((_batch_size, opt.latent_dim))

    # Generate a batch of images
    _gen_imgs1, _gen_imgs2 = coupled_generators(z)
    # Determine validity of generated images
    validity1, validity2 = coupled_discriminators(_gen_imgs1, _gen_imgs2)

    _g_loss = (adversarial_loss(validity1, _valid) + adversarial_loss(validity2, _valid)) / 2

    return _g_loss, _gen_imgs1, _gen_imgs2


def d_forward(_imgs1, _imgs2, _gen_imgs1, _gen_imgs2, _valid, _fake):
    """Discriminator forward function"""
    # Determine validity of real and generated images
    validity1_real, validity2_real = coupled_discriminators(_imgs1, _imgs2)
    validity1_fake, validity2_fake = coupled_discriminators(_gen_imgs1, _gen_imgs2)

    _d_loss = (
                      adversarial_loss(validity1_real, _valid)
                      + adversarial_loss(validity1_fake, _fake)
                      + adversarial_loss(validity2_real, _valid)
                      + adversarial_loss(validity2_fake, _fake)
              ) / 4

    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

coupled_generators.set_train()
coupled_discriminators.set_train()

# ----------
#  Training
# ----------

for epoch in range(opt.n_epochs):
    for i, ((imgs1, _), (imgs2, _)) in enumerate(zip(dataset1.create_tuple_iterator(),
                                                     dataset2.create_tuple_iterator())):
        batch_size = imgs1.shape[0]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((batch_size, 1)))
        fake = ops.stop_gradient(ops.zeros((batch_size, 1)))

        # Configure input
        imgs1 = imgs1.broadcast_to((imgs1.shape[0], 3, opt.img_size, opt.img_size))

        # ------------------
        #  Train Generators
        # ------------------

        (g_loss, gen_imgs1, gen_imgs2), g_grads = grad_g(batch_size, valid)
        optimizer_G(g_grads)

        # ----------------------
        #  Train Discriminators
        # ----------------------

        (d_loss), d_grads = grad_d(imgs1, imgs2, ops.stop_gradient(gen_imgs1),
                                   ops.stop_gradient(gen_imgs2), valid, fake)
        optimizer_D(d_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset1.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset1.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            gen_imgs = ops.cat((gen_imgs1, gen_imgs2), 0)
            to_image(gen_imgs, os.path.join("images", F'{batches_done}.png'))

implementations/cogan/cogan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/cogan/img_utils.py
mindspore-courses/MindSpore-GAN
"""Dataset setting and data loader for MNIST-M.

Modified from
https://github.com/pytorch/vision/blob/master/torchvision/datasets/mnist.py

CREDIT: https://github.com/corenel
"""

from __future__ import print_function

import errno
import os
import gzip
import pickle

import mindspore
from PIL import Image
from mindspore import Tensor
import mindspore.common.dtype as mstype

# import essential packages
from six.moves import urllib


class MNISTM:
    """`MNIST-M Dataset."""

    url = "https://github.com/VanushVaswani/keras_mnistm/releases/download/1.0/keras_mnistm.pkl.gz"

    raw_folder = "raw"
    processed_folder = "processed"
    training_file = "mnist_m_train.pt"
    test_file = "mnist_m_test.pt"

    mnist_train_labels = []
    mnist_test_labels = []

    def __init__(self, root, mnist_root="data", train=True, transform=None, target_transform=None):
        """Init MNIST-M dataset."""
        super().__init__(MNISTM)
        self.root = os.path.expanduser(root)
        self.mnist_root = os.path.expanduser(mnist_root)
        self.transform = transform
        self.target_transform = target_transform
        self.train = train  # training set or test set

        self.train_ds, self.test_ds = self.dl_preprocess()

    def __getitem__(self, index):
        """Get images and target for data loader.

        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        if self.train:
            img, target = self.train_ds[0][index], self.train_ds[1][index]
        else:
            img, target = self.test_ds[0][index], self.test_ds[1][index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.squeeze(), mode="RGB")

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img[0], target

    def __len__(self):
        """Return size of dataset."""
        if self.train:
            return len(self.train_ds[0])
        return len(self.test_ds[0])

    def get_train_labels(self, labels):
        """Get MNIST labels"""
        for _, label in enumerate(labels.create_tuple_iterator()):
            self.mnist_train_labels.append(label)

    def get_test_labels(self, labels):
        """Get MNIST labels"""
        for _, label in enumerate(labels.create_tuple_iterator()):
            self.mnist_test_labels.append(label)

    def dl_preprocess(self):
        """Create MNIST-M dataset."""

        # make data dirs
        try:
            os.makedirs(os.path.join(self.root, self.raw_folder))
        except OSError as e:
            if e.errno == errno.EEXIST:
                pass
            else:
                raise

        # download pkl files
        filename = self.url.rpartition("/")[2]
        file_path = os.path.join(self.root, self.raw_folder, filename)
        if not os.path.exists(os.path.join(self.root, 'raw/keras_mnistm.pkl')):
            print("Downloading " + self.url)
            if not os.path.exists(file_path.replace(".gz", "")):
                data = urllib.request.urlopen(self.url)
                with open(file_path, "wb") as f:
                    f.write(data.read())
                with open(file_path.replace(".gz", ""), "wb") as out_f, gzip.GzipFile(file_path) as zip_f:
                    out_f.write(zip_f.read())
                os.unlink(file_path)

        # process and save as torch files
        print("Processing...")

        # load MNIST-M images from pkl file
        with open(file_path.replace(".gz", ""), "rb") as f:
            mnist_m_data = pickle.load(f, encoding="bytes")

        mnist_m_train_data = Tensor(mnist_m_data[b"train"], dtype=mstype.byte)
        mnist_m_test_data = Tensor(mnist_m_data[b"test"], dtype=mstype.byte)

        # get MNIST labels
        mnist_train = mindspore.dataset.MnistDataset(
            dataset_dir=self.mnist_root,
            usage='train',
            shuffle=False)
        mnist_test = mindspore.dataset.MnistDataset(
            dataset_dir=self.mnist_root,
            usage='test',
            shuffle=False)

        for _, (_, label) in enumerate(mnist_train.create_tuple_iterator()):
            self.mnist_train_labels.append(label)

        for _, (_, label) in enumerate(mnist_test.create_tuple_iterator()):
            self.mnist_test_labels.append(label)

        # save MNIST-M dataset
        training_set = (mnist_m_train_data.asnumpy(), self.mnist_train_labels)
        test_set = (mnist_m_test_data.asnumpy(), self.mnist_test_labels)

        print("Done!")
        return training_set, test_set

implementations/cogan/mnistm.py
mindspore-courses/MindSpore-GAN
"""
Inpainting using Generative Adversarial Networks.
The dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0
(if not available there see if options are listed at http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)
Instrustion on running the script:
1. Download the dataset from the provided link
2. Save the folder 'img_align_celeba' to '../../data/'
4. Run the sript using command 'python3 context_encoder.py'
"""

import argparse
import os

import numpy as np
from mindspore import nn
from mindspore import ops
from mindspore.dataset import CelebADataset
from mindspore.dataset.vision import transforms, Inter

from img_utils import to_image
from models import Generator, Discriminator

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=8, help="size of the batches")
parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=4, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=128, help="size of each image dimension")
parser.add_argument("--mask_size", type=int, default=64, help="size of random mask")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=500, help="interval between image sampling")
opt = parser.parse_args()
print(opt)

# Calculate output of image discriminator (PatchGAN)
patch_h, patch_w = int(opt.mask_size / 2 ** 3), int(opt.mask_size / 2 ** 3)
patch = (1, patch_h, patch_w)

adversarial_loss = nn.MSELoss()
pixelwise_loss = nn.L1Loss()

# Initialize generator and discriminator
generator = Generator(channels=opt.channels)
discriminator = Discriminator(channels=opt.channels)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), opt.lr, opt.b1, opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), opt.lr, opt.b1, opt.b2)


def apply_random_mask(img, img_size, mask_size):
    """Randomly masks image"""
    y1, x1 = np.random.randint(0, img_size - mask_size, 2)
    y2, x2 = y1 + mask_size, x1 + mask_size
    _masked_part = img[:, y1:y2, x1:x2]
    _masked_img = img.copy()
    _masked_img[:, y1:y2, x1:x2] = 1

    return _masked_img, _masked_part


def apply_center_mask(img, img_size, mask_size):
    """Mask center part of image"""
    # Get upper-left pixel coordinate
    k = (img_size - mask_size) // 2
    masked_img = img.copy()
    masked_img[:, k: k + mask_size, k: k + mask_size] = 1

    return masked_img, k


def train_preprocess(_imgs):
    """Dataset preprocess func"""
    _imgs = transforms.Resize((opt.img_size, opt.img_size), Inter.BICUBIC)(_imgs)
    _imgs = transforms.ToTensor()(_imgs)
    _imgs = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False)(_imgs)
    # For training data perform random mask
    _masked_img, _aux = apply_random_mask(_imgs, 128, 64)
    return _imgs, _masked_img, _aux


def test_preprocess(_imgs):
    """Dataset preprocess func"""
    _imgs = transforms.Resize((opt.img_size, opt.img_size), Inter.BICUBIC)(_imgs)
    _imgs = transforms.ToTensor()(_imgs)
    _imgs = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False)(_imgs)
    # For test data mask the center of the image
    _masked_img, _aux = apply_center_mask(_imgs, 128, 64)
    return _imgs, _masked_img, _aux


train_dataset = CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True
).map(operations=[train_preprocess], input_columns=["image"], output_columns=["image_hr", "masked_image", "aux"]).batch(
    opt.batch_size)

test_dataset = CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True
).map(operations=[test_preprocess], input_columns=["image"], output_columns=["image_hr", "masked_image", "aux"]).batch(
    12)


def save_sample(_batches_done):
    """Save sample images"""
    samples, masked_samples, j, _ = next(test_dataset.create_tuple_iterator())
    j = j[0].asnumpy().item()  # Upper-left coordinate of mask
    # Generate inpainted image
    gen_mask = generator(masked_samples)
    filled_samples = masked_samples.copy()
    filled_samples[:, :, j: j + opt.mask_size, j: j + opt.mask_size] = gen_mask
    # Save sample
    sample = ops.cat((masked_samples, filled_samples, samples), -2)
    to_image(sample, f'images/{_batches_done}.png')


def g_forward(_masked_imgs, _masked_parts, _valid):
    """Generator warmup forward func"""
    # Generate a batch of images
    _gen_parts = generator(_masked_imgs)

    # Adversarial and pixelwise loss
    _g_adv = adversarial_loss(discriminator(_gen_parts), _valid)
    _g_pixel = pixelwise_loss(_gen_parts, _masked_parts)
    # Total loss
    _g_loss = 0.001 * _g_adv + 0.999 * _g_pixel
    return _g_loss, _g_adv, _g_pixel, _gen_parts


def d_forward(_masked_parts, _gen_parts, _valid, _fake):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_loss = adversarial_loss(discriminator(_masked_parts), _valid)
    fake_loss = adversarial_loss(discriminator(_gen_parts), _fake)
    _d_loss = 0.5 * (real_loss + fake_loss)
    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

for epoch in range(opt.n_epochs):
    # Model inputs
    for i, (imgs, masked_imgs, masked_parts, _) in enumerate(train_dataset.create_tuple_iterator()):
        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs.shape[0], *patch)))
        fake = ops.stop_gradient(ops.zeros((imgs.shape[0], *patch)))

        # -----------------
        #  Train Generator
        # -----------------
        (g_loss, g_adv, g_pixel, gen_parts), g_grads = grad_g(masked_imgs, masked_parts, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (d_loss), d_grads = grad_d(masked_parts, ops.stop_gradient(gen_parts), valid, fake)
        optimizer_D(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{train_dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G adv: {g_adv.asnumpy().item():.4f}, '
            f'pixel: {g_pixel.asnumpy().item():.4f}]'
        )

        # Generate sample at sample interval
        batches_done = epoch * train_dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            save_sample(batches_done)

implementations/context_encoder/context_encoder.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 6,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/context_encoder/img_utils.py
mindspore-courses/MindSpore-GAN
"""The Models of Context Encoder"""

import mindspore.common.initializer as init
from mindspore import nn


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self, channels=3):
        super().__init__(Generator)

        def downsample(in_feat, out_feat, normalize=True):
            layers = [nn.Conv2d(in_feat, out_feat, 4, 2, 'pad', 1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.BatchNorm2d(out_feat, 0.8,
                                             gamma_init=init.Normal(0.02, 1.0),
                                             beta_init=init.Constant(0.0), affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        def upsample(in_feat, out_feat, normalize=True):
            layers = [nn.Conv2dTranspose(in_feat, out_feat, 4, 2, 'pad', 1,
                                         weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.BatchNorm2d(out_feat, 0.8,
                                             gamma_init=init.Normal(0.02, 1.0),
                                             beta_init=init.Constant(0.0), affine=False))
            layers.append(nn.ReLU())
            return layers

        self.model = nn.SequentialCell(
            *downsample(channels, 64, normalize=False),
            *downsample(64, 64),
            *downsample(64, 128),
            *downsample(128, 256),
            *downsample(256, 512),
            nn.Conv2d(512, 4000, 1,
                      weight_init=init.Normal(0.02, 0.0)),
            *upsample(4000, 512),
            *upsample(512, 256),
            *upsample(256, 128),
            *upsample(128, 64),
            nn.Conv2d(64, channels, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, x):
        return self.model(x)


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, channels=3):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, stride, normalize):
            """Returns layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 3, stride, 'pad', 1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.InstanceNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        layers = []
        in_filters = channels
        for out_filters, stride, normalize in [(64, 2, False), (128, 2, True), (256, 2, True), (512, 1, True)]:
            layers.extend(discriminator_block(in_filters, out_filters, stride, normalize))
            in_filters = out_filters

        layers.append(nn.Conv2d(out_filters, 1, 3, 1, 'pad', 1,
                                weight_init=init.Normal(0.02, 0.0)))

        self.model = nn.SequentialCell(*layers)

    def construct(self, img):
        return self.model(img)

implementations/context_encoder/models.py
mindspore-courses/MindSpore-GAN
"""CycleGAN Model"""

import argparse
import datetime
import os
import sys
import time
import urllib.request
import zipfile

import mindspore
from mindspore import nn
from mindspore import ops
from mindspore.dataset import MindDataset

from img_utils import to_image, make_grid
from models import GeneratorResNet, Discriminator
from utils import DynamicDecayLR, ReplayBuffer

file_path = "../../data/"

if not os.path.exists(file_path + "apple2orange"):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path + "apple2orange")
    url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/models/application/CycleGAN_apple2orange.zip'
    print("Downloading apple2orange dataset from " + url)
    urllib.request.urlretrieve(url, os.path.join(file_path, 'apple2orange.zip'))

    while not os.path.exists(os.path.join(file_path, 'apple2orange.zip')):
        pass

    zippedData = zipfile.ZipFile(os.path.join(file_path, 'apple2orange.zip'), 'r')
    print("Unzipping dataset...")
    zippedData.extractall(path=file_path)
    zippedData.close()
    os.rename(file_path + 'CycleGAN_apple2orange', file_path + 'apple2orange')
    os.remove(os.path.join(file_path, 'apple2orange.zip'))

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="apple2orange", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=256, help="size of image height")
parser.add_argument("--img_width", type=int, default=256, help="size of image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=1, help="interval between saving generator outputs")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between saving model checkpoints")
parser.add_argument("--n_residual_blocks", type=int, default=9, help="number of residual blocks in generator")
parser.add_argument("--lambda_cyc", type=float, default=10.0, help="cycle loss weight")
parser.add_argument("--lambda_id", type=float, default=5.0, help="identity loss weight")
opt = parser.parse_args()
print(opt)

# Create sample and checkpoint directories
os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

# Losses
criterion_GAN = nn.MSELoss()
criterion_cycle = nn.L1Loss()
criterion_identity = nn.L1Loss()

input_shape = (opt.channels, opt.img_height, opt.img_width)

# Initialize generator and discriminator
G_AB = GeneratorResNet(input_shape, opt.n_residual_blocks)
G_BA = GeneratorResNet(input_shape, opt.n_residual_blocks)
D_A = Discriminator(input_shape)
D_B = Discriminator(input_shape)

G_AB.update_parameters_name("G_AB")
G_BA.update_parameters_name("G_BA")
D_A.update_parameters_name("G_AB")
D_B.update_parameters_name("G_BA")

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G_AB_{opt.epoch}.ckpt', G_AB)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G_BA_{opt.epoch}.ckpt', G_BA)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_A_{opt.epoch}.ckpt', D_A)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_B_{opt.epoch}.ckpt', D_B)


# Learning rate update schedulers
decay_lr_G_AB = DynamicDecayLR(opt.lr, opt.n_epochs, 1019, opt.epoch, opt.decay_epoch)
decay_lr_G_BA = DynamicDecayLR(opt.lr, opt.n_epochs, 1019, opt.epoch, opt.decay_epoch)
decay_lr_D_A = DynamicDecayLR(opt.lr, opt.n_epochs, 1019, opt.epoch, opt.decay_epoch)
decay_lr_D_B = DynamicDecayLR(opt.lr, opt.n_epochs, 1019, opt.epoch, opt.decay_epoch)

optimizer_G_AB = nn.optim.Adam(G_AB.trainable_params(), learning_rate=decay_lr_G_AB,
                               beta1=opt.b1, beta2=opt.b2)
optimizer_G_BA = nn.optim.Adam(G_BA.trainable_params(), learning_rate=decay_lr_G_BA,
                               beta1=opt.b1, beta2=opt.b2)
optimizer_D_A = nn.optim.Adam(D_A.trainable_params(), learning_rate=decay_lr_D_A,
                              beta1=opt.b1, beta2=opt.b2)

optimizer_D_B = nn.optim.Adam(D_B.trainable_params(), learning_rate=decay_lr_D_B,
                              beta1=opt.b1, beta2=opt.b2)

optimizer_G_AB.update_parameters_name("optimizer_G_AB")
optimizer_G_BA.update_parameters_name("optimizer_G_BA")
optimizer_D_A.update_parameters_name("optimizer_D_A")
optimizer_D_B.update_parameters_name("optimizer_D_B")

# Buffers of previously generated samples
fake_A_buffer = ReplayBuffer()
fake_B_buffer = ReplayBuffer()

dataset = MindDataset("../../data/apple2orange/apple2orange_train.mindrecord", shuffle=True).batch(opt.batch_size)

val_dataset = MindDataset("../../data/apple2orange/apple2orange_train.mindrecord", shuffle=True).batch(5)


def sample_images(batches):
    """Saves a generated sample from the validation set"""
    imgs = next(val_dataset.create_dict_iterator())
    G_AB.set_train(False)
    G_BA.set_train(False)
    _real_A = imgs["image_A"]
    _fake_B = G_AB(_real_A)
    _real_B = imgs["image_B"]
    _fake_A = G_BA(_real_B)
    # Arange images along x-axis
    _real_A = make_grid(_real_A, nrow=5, normalize=True)
    _real_B = make_grid(_real_B, nrow=5, normalize=True)
    _fake_A = make_grid(_fake_A, nrow=5, normalize=True)
    _fake_B = make_grid(_fake_B, nrow=5, normalize=True)
    # Arange images along y-axis
    image_grid = ops.cat((_real_A, _fake_B, _real_B, _fake_A), 1)
    to_image(image_grid, os.path.join(f'images/{opt.dataset_name}', F'{batches}.png'))


def g_forward(_real_A, _real_B, _valid):
    """Generator forward function"""
    _fake_A = G_BA(_real_B)
    _fake_B = G_AB(_real_A)
    _rec_A = G_BA(_fake_B)
    _rec_B = G_AB(_fake_A)
    _identity_A = G_BA(_real_A)
    _identity_B = G_AB(_real_B)
    _loss_G_A = criterion_GAN(D_B(_fake_B), _valid)
    _loss_G_B = criterion_GAN(D_A(_fake_A), _valid)
    _loss_cycle_A = criterion_cycle(_rec_A, _real_A) * opt.lambda_cyc
    _loss_cycle_B = criterion_cycle(_rec_B, _real_B) * opt.lambda_cyc
    _loss_idt_A = criterion_identity(_identity_A, _real_A) * opt.lambda_cyc * opt.lambda_id
    _loss_idt_B = criterion_identity(_identity_B, _real_B) * opt.lambda_cyc * opt.lambda_id
    loss_G_ = _loss_G_A + _loss_G_B + _loss_cycle_A + _loss_cycle_B + _loss_idt_A + _loss_idt_B

    return loss_G_, _loss_G_A, _loss_G_B, _loss_cycle_A, _loss_cycle_B, _loss_idt_A, _loss_idt_B, _fake_A, _fake_B


def d_a_forward(_real_A, _fake_A, _valid, _fake):
    """Discriminator forward function"""
    # Real loss
    loss_real_ = criterion_GAN(D_A(_real_A), _valid)
    # Fake loss (on batch of previously generated samples)
    _fake_A_ = fake_A_buffer.push_and_pop(_fake_A)
    loss_fake_ = criterion_GAN(D_A(ops.stop_gradient(_fake_A_)), _fake)
    # Total loss
    loss_D_A_ = (loss_real_ + loss_fake_) / 2

    return loss_D_A_


def d_b_forward(_real_B, _fake_B, _valid, _fake):
    """Discriminator forward function"""
    # Real loss
    loss_real_ = criterion_GAN(D_B(_real_B), _valid)
    # Fake loss (on batch of previously generated samples)
    _fake_B_ = fake_B_buffer.push_and_pop(_fake_B)
    loss_fake_ = criterion_GAN(D_B(ops.stop_gradient(_fake_B_)), _fake)
    # Total loss
    loss_D_B_ = (loss_real_ + loss_fake_) / 2

    return loss_D_B_


grad_g_ab = ops.value_and_grad(g_forward, None, optimizer_G_AB.parameters)
grad_g_ba = ops.value_and_grad(g_forward, None, optimizer_G_BA.parameters)
grad_d_a = ops.value_and_grad(d_a_forward, None, optimizer_D_A.parameters)
grad_d_b = ops.value_and_grad(d_b_forward, None, optimizer_D_B.parameters)

# ----------
#  Training
# ----------

G_AB.set_train()
G_BA.set_train()
D_A.set_train()
D_B.set_train()

prev_time = time.time()
for epoch in range(opt.n_epochs):
    # Model inputs
    for i, batch in enumerate(dataset.create_dict_iterator()):
        # Set model input
        real_A = batch["image_A"]
        real_B = batch["image_B"]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((real_A.shape[0], *D_A.output_shape)))
        fake = ops.stop_gradient(ops.zeros((real_A.shape[0], *D_A.output_shape)))

        # ------------------
        #  Train Generators
        # ------------------

        loss_G, loss_G_A, loss_G_B, loss_cycle_A, loss_cycle_B, loss_idt_A, loss_idt_B, fake_A, fake_B = g_forward(
            real_A, real_B, valid)

        (_, _, _, _, _, _, _, _, _), g_ab_grads = grad_g_ab(real_A, real_B, valid)
        (_, _, _, _, _, _, _, _, _), g_ba_grads = grad_g_ba(real_A, real_B, valid)
        optimizer_G_AB(g_ab_grads)
        optimizer_G_BA(g_ba_grads)

        # -----------------------
        #  Train Discriminator A
        # -----------------------

        (loss_D_A), d_a_grads = grad_d_a(real_A, fake_A, valid, fake)
        optimizer_D_A(d_a_grads)

        # -----------------------
        #  Train Discriminator B
        # -----------------------

        (loss_D_B), d_b_grads = grad_d_b(real_B, fake_B, valid, fake)
        optimizer_D_B(d_b_grads)

        loss_D = (loss_D_A + loss_D_B) / 2

        # --------------
        #  Log Progress
        # --------------

        # Determine approximate time left
        batches_done = epoch * dataset.get_dataset_size() + i
        batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {loss_D.asnumpy().item():.4f}] '
            f'[G loss: {loss_G.asnumpy().item():.4f}, '
            f'adv: A: {loss_G_A.asnumpy().item():.4f} B: {loss_G_B.asnumpy().item():.4f}, '
            f'cycle: A: {loss_cycle_A.asnumpy().item():.4f}, B: {loss_cycle_B.asnumpy().item():.4f}, '
            f'identity: A: {loss_idt_A.asnumpy().item():.4f}, B: {loss_idt_B.asnumpy().item():4f}] '
            f'ETA: {time_left}'
        )

        # If at sample interval save image
        if batches_done % opt.sample_interval == 0:
            sample_images(batches_done)

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        mindspore.save_checkpoint(G_AB, f'saved_models/{opt.dataset_name}/G_AB_{epoch}.ckpt')
        mindspore.save_checkpoint(G_BA, f'saved_models/{opt.dataset_name}/G_BA_{epoch}.ckpt')
        mindspore.save_checkpoint(D_A, f'saved_models/{opt.dataset_name}/D_A_{epoch}.ckpt')
        mindspore.save_checkpoint(D_B, f'saved_models/{opt.dataset_name}/D_B_{epoch}.ckpt')

implementations/cyclegan/cyclegan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/cyclegan/img_utils.py
mindspore-courses/MindSpore-GAN
"""CycleGAN Networks"""

from mindspore import nn
import mindspore.common.initializer as init


##############################
#           RESNET
##############################

class ResidualBlock(nn.Cell):
    """Residual Block"""
    def __init__(self, in_features):
        super().__init__(ResidualBlock)

        self.block = nn.SequentialCell(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3, has_bias=False,
                      pad_mode='valid',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(in_features, affine=False),
            nn.ReLU(),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3, has_bias=False,
                      pad_mode='valid',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(in_features, affine=False)
        )

    def construct(self, x):
        return x + self.block(x)


class GeneratorResNet(nn.Cell):
    """Generator Network"""
    def __init__(self, input_shape, num_residual_blocks):
        super().__init__(GeneratorResNet)

        channels = input_shape[0]

        # Initial convolution block
        out_features = 64
        model = [
            nn.ReflectionPad2d(channels),
            nn.Conv2d(channels, out_features, 7, has_bias=False,
                      pad_mode='valid',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(out_features, affine=False),
            nn.ReLU()
        ]
        in_features = out_features

        # Downsampling
        for _ in range(2):
            out_features *= 2
            model += [
                nn.Conv2d(in_features, out_features, 3, stride=2, has_bias=False,
                          pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.BatchNorm2d(out_features, affine=False),
                nn.ReLU()
            ]
            in_features = out_features

        # Residual blocks
        for _ in range(num_residual_blocks):
            model += [ResidualBlock(out_features)]

        # Upsampling
        for _ in range(2):
            out_features //= 2
            model += [
                nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
                nn.Conv2d(in_features, out_features, 3, has_bias=False,
                          stride=1, pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.BatchNorm2d(out_features, affine=False),
                nn.ReLU()
            ]
            in_features = out_features

        # Output layer
        model += [
            nn.ReflectionPad2d(channels),
            nn.Conv2d(out_features, channels, 7, has_bias=False,
                      pad_mode='valid',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        ]

        self.model = nn.SequentialCell(*model)

    def construct(self, x):
        return self.model(x)

    ##############################
    #        Discriminator
    ##############################


class Discriminator(nn.Cell):
    """Discriminator network"""
    def __init__(self, input_shape):
        super().__init__(Discriminator)

        channels, height, width = input_shape

        # Calculate output shape of image discriminator (PatchGAN)
        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)

        def discriminator_block(in_filters, out_filters, normalize=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, has_bias=False,
                                stride=2, pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *discriminator_block(channels, 64, normalize=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            *discriminator_block(256, 512),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, 4, has_bias=False,
                      pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0))
        )

    def construct(self, img):
        return self.model(img)

implementations/cyclegan/models.py
mindspore-courses/MindSpore-GAN
"""Utilities"""


import random

from mindspore import ops, Tensor
from mindspore.nn.learning_rate_schedule import LearningRateSchedule
from mindspore.common import dtype as mstype


class DynamicDecayLR(LearningRateSchedule):
    """Learning rate schedule for dynamic"""

    def __init__(self, lr, n_epochs, step_per_epoch, offset, decay_start_epoch):
        super().__init__()
        assert (n_epochs - decay_start_epoch) > 0, "Decay must start before the training session ends!"
        self.lr = lr
        self.n_epochs = n_epochs
        self.step_per_epoch = step_per_epoch
        self.offset = offset
        self.decay_start_epoch = decay_start_epoch
        self.cast = ops.Cast()

    def construct(self, global_step):
        epoch = self.cast(global_step, mstype.float32) // self.step_per_epoch
        return self.lr * (1.0 - max(Tensor(0.0), epoch + self.offset - self.decay_start_epoch) / (
                self.n_epochs - self.decay_start_epoch))


class ReplayBuffer:
    """Replay buffer"""

    def __init__(self, max_size=50):
        assert max_size > 0, "Empty buffer or trying to create a black hole. Be careful."
        self.max_size = max_size
        self.data = []

    def push_and_pop(self, data):
        """Push and pop function"""
        to_return = []
        for element in data:
            element = ops.unsqueeze(element, 0)
            if len(self.data) < self.max_size:
                self.data.append(element)
                to_return.append(element)
            else:
                if random.uniform(0, 1) > 0.5:
                    i = random.randint(0, self.max_size - 1)
                    to_return.append(self.data[i].copy())
                    self.data[i] = element
                else:
                    to_return.append(element)
        return ops.cat(to_return)

implementations/cyclegan/utils.py
mindspore-courses/MindSpore-GAN
"""DCGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size",
                    type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.init_size = opt.img_size // 4
        self.l1 = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 128 * self.init_size ** 2)
        )

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, 1, 'pad', 1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, _z):
        out = self.l1(_z)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [
                nn.Conv2d(in_filters, out_filters, 3, 2, 'pad', 1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0)))
            return block

        self.model = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128)
        )

        ds_size = opt.img_size // 2 ** 4
        self.adv_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, 1),
            nn.Sigmoid()
        )

    def construct(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        return validity


# Loss function
adversarial_loss = nn.BCELoss()

generator = Generator()
discriminator = Discriminator()

G_Optim = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
D_Optim = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_z, _valid):
    """Generator forward function"""
    _gen_imgs = generator(_z)
    _g_loss = adversarial_loss(discriminator(_gen_imgs), _valid)
    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""
    real_loss = adversarial_loss(discriminator(_real_imgs), _valid)
    fake_loss = adversarial_loss(discriminator(_gen_imgs), _fake)
    _d_loss = (real_loss + fake_loss) / 2

    return _d_loss


transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, G_Optim.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, D_Optim.parameters, has_aux=False)

for epoch in range(opt.n_epochs):
    generator.set_train()
    discriminator.set_train()
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        valid = ops.ones((imgs.shape[0], 1))
        fake = ops.zeros((imgs.shape[0], 1))
        ops.stop_gradient(valid)
        ops.stop_gradient(fake)

        real_imgs = Tensor(imgs, dtype=mstype.float32)

        z = ops.randn((imgs.shape[0], opt.latent_dim))

        (g_loss, gen_imgs), g_grads = grad_g(z, valid)
        G_Optim(g_grads)

        (d_loss), d_grads = grad_d(real_imgs, gen_imgs, valid, fake)
        D_Optim(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item()}] [G loss: {g_loss.asnumpy().item()}]'
        )

        batches_done = epoch * dataset.get_dataset_size() + i

        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/dcgan/dcgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/dcgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""Dataset Preprocess"""

import glob
import os

import numpy as np
from PIL import Image
from mindspore.dataset.transforms import Compose


class Edges2ShoesDataset():
    """Edges2Shoes Dataset"""

    def __init__(self, root, transforms_=None, mode='train'):
        self.transform = Compose(transforms_)

        self.files = sorted(glob.glob(os.path.join(root, mode) + '/*.*'))

    def __getitem__(self, index):
        img = Image.open(self.files[index % len(self.files)])
        w, h = img.size
        img_A = img.crop((0, 0, w / 2, h))
        img_B = img.crop((w / 2, 0, w, h))

        if np.random.random() < 0.5:
            img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], "RGB")
            img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], "RGB")

        img_A = self.transform(img_A)
        img_B = self.transform(img_B)

        return img_A[0], img_B[0]

    def __len__(self):
        return len(self.files)

implementations/discogan/datasets.py
mindspore-courses/MindSpore-GAN
"""DiscoGAN Model"""

import argparse
import datetime
import itertools
import os
import sys
import time

import mindspore
from mindspore import nn
from mindspore import ops
from mindspore.dataset.vision import transforms, Inter

from datasets import Edges2ShoesDataset
from img_utils import to_image
from models import GeneratorUNet, Discriminator

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="edges2shoes", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=64, help="size of image height")
parser.add_argument("--img_width", type=int, default=64, help="size of image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=100, help="interval between saving generator samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
opt = parser.parse_args()
print(opt)

# Create sample and checkpoint directories
os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

# Losses
adversarial_loss = nn.MSELoss()
cycle_loss = nn.L1Loss()
pixelwise_loss = nn.L1Loss()

input_shape = (opt.channels, opt.img_height, opt.img_width)

# Initialize generator and discriminator
G_AB = GeneratorUNet(input_shape)
G_BA = GeneratorUNet(input_shape)
D_A = Discriminator(input_shape)
D_B = Discriminator(input_shape)

G_AB.update_parameters_name("G_AB")
G_BA.update_parameters_name("G_BA")
D_A.update_parameters_name("G_AB")
D_B.update_parameters_name("G_BA")

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G_AB_{opt.epoch}.ckpt', G_AB)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G_BA_{opt.epoch}.ckpt', G_BA)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_A_{opt.epoch}.ckpt', D_A)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_B_{opt.epoch}.ckpt', D_B)

# Optimizers
optimizer_G = nn.optim.Adam(itertools.chain(G_AB.trainable_params(), G_BA.trainable_params()),
                            learning_rate=opt.lr,
                            beta1=opt.b1, beta2=opt.b2)
optimizer_D_A = nn.optim.Adam(D_A.trainable_params(), learning_rate=opt.lr,
                              beta1=opt.b1, beta2=opt.b2)

optimizer_D_B = nn.optim.Adam(D_B.trainable_params(), learning_rate=opt.lr,
                              beta1=opt.b1, beta2=opt.b2)

# Dataset loader
transforms_ = [
    transforms.Resize((opt.img_height, opt.img_width), Inter.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
]

dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        transforms_=transforms_
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(opt.batch_size)

val_dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        transforms_=transforms_,
        mode='val'
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(16)


def sample_image(batches):
    """Saves a generated sample from the validation set"""
    imgs = next(val_dataset.create_tuple_iterator())
    G_AB.set_train(False)
    G_BA.set_train(False)
    _real_A = imgs[0]
    _fake_B = G_AB(_real_A)
    _real_B = imgs[1]
    _fake_A = G_BA(_real_B)
    img_sample = ops.cat((_real_A, _fake_B, _real_B, _fake_A), 0)
    to_image(img_sample, os.path.join(f'images/{opt.dataset_name}', f'{batches}.png'))


def g_forward(_real_A, _real_B, _valid):
    """Generator warmup forward func"""
    # GAN loss
    _fake_B = G_AB(_real_A)
    loss_GAN_AB = adversarial_loss(D_B(_fake_B), _valid)
    _fake_A = G_BA(_real_B)
    loss_GAN_BA = adversarial_loss(D_A(_fake_A), _valid)

    _loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2

    # Pixelwise translation loss
    _loss_pixelwise = (pixelwise_loss(_fake_A, _real_A) +
                       pixelwise_loss(_fake_B, _real_B)) / 2

    # Cycle loss
    loss_cycle_A = cycle_loss(G_BA(_fake_B), _real_A)
    loss_cycle_B = cycle_loss(G_AB(_fake_A), _real_B)
    _loss_cycle = (loss_cycle_A + loss_cycle_B) / 2

    # Total loss
    _loss_G = _loss_GAN + _loss_cycle + _loss_pixelwise

    return _loss_G, _loss_GAN, _loss_pixelwise, _loss_cycle, _fake_A, _fake_B


def d_a_forward(_real_A, _fake_A, _valid, _fake):
    """Discriminator A forward function"""

    # Real loss
    loss_real = adversarial_loss(D_A(_real_A), _valid)
    # Fake loss (on batch of previously generated samples)
    loss_fake = adversarial_loss(D_A(_fake_A), _fake)
    # Total loss
    _loss_D_A = (loss_real + loss_fake) / 2

    return _loss_D_A


def d_b_forward(_real_B, _fake_B, _valid, _fake):
    """Discriminator B forward function"""

    # Real loss
    loss_real = adversarial_loss(D_A(_real_B), _valid)
    # Fake loss (on batch of previously generated samples)
    loss_fake = adversarial_loss(D_A(_fake_B), _fake)
    # Total loss
    _loss_D_B = (loss_real + loss_fake) / 2

    return _loss_D_B


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d_a = ops.value_and_grad(d_a_forward, None, optimizer_D_A.parameters, has_aux=False)
grad_d_b = ops.value_and_grad(d_b_forward, None, optimizer_D_B.parameters, has_aux=False)

# ----------
#  Training
# ----------

prev_time = time.time()
for epoch in range(opt.n_epochs):
    # Model inputs
    for i, (real_A, real_B) in enumerate(dataset.create_tuple_iterator()):

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((real_A.shape[0], *D_A.output_shape)))
        fake = ops.stop_gradient(ops.zeros((real_A.shape[0], *D_A.output_shape)))

        # ------------------
        #  Train Generators
        # ------------------

        G_AB.set_train()
        G_BA.set_train()

        (loss_G, loss_GAN, loss_pixelwise, loss_cycle, fake_A, fake_B), g_grads = grad_g(real_A, real_B, valid)
        optimizer_G(g_grads)

        # -----------------------
        #  Train Discriminator A
        # -----------------------

        D_A.set_train()

        (loss_D_A), d_a_grads = grad_d_a(real_A, ops.stop_gradient(fake_A), valid, fake)
        optimizer_D_A(d_a_grads)

        # -----------------------
        #  Train Discriminator B
        # -----------------------

        D_B.set_train()

        (loss_D_B), d_b_grads = grad_d_b(real_B, ops.stop_gradient(fake_B), valid, fake)
        optimizer_D_B(d_b_grads)

        loss_D = 0.5 * (loss_D_A + loss_D_B)

        # --------------
        #  Log Progress
        # --------------

        # Determine approximate time left
        batches_done = epoch * dataset.get_dataset_size() + i
        batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D VAE loss: {loss_D.asnumpy().item():.4f}] '
            f'[G loss: {loss_G.asnumpy().item():.4f}, '
            f'adv: {loss_GAN.asnumpy().item():.4f}, '
            f'pixel: {loss_pixelwise.asnumpy().item():.4f}, '
            f'cycle:{loss_cycle.asnumpy().item():.4f}] '
            f'ETA: {time_left}'
        )

        if batches_done % opt.sample_interval == 0:
            sample_image(batches_done)

        if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
            # Save model checkpoints
            mindspore.save_checkpoint(G_AB, f'saved_models/{opt.dataset_name}/G_AB_{epoch}.ckpt')
            mindspore.save_checkpoint(G_BA, f'saved_models/{opt.dataset_name}/G_BA_{epoch}.ckpt')
            mindspore.save_checkpoint(D_A, f'saved_models/{opt.dataset_name}/D_A_{epoch}.ckpt')
            mindspore.save_checkpoint(D_B, f'saved_models/{opt.dataset_name}/D_B_{epoch}.ckpt')

implementations/discogan/discogan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/discogan/img_utils.py
mindspore-courses/MindSpore-GAN
"""DiscoGAN Networks"""

import mindspore.common.initializer as init
from mindspore import nn, ops


##############################
#           U-NET
##############################


class UNetDown(nn.Cell):
    """U-Net Down"""

    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
        super().__init__(UNetDown)
        layers = [nn.Conv2d(in_size, out_size, 4, stride=2, pad_mode='pad', padding=1, has_bias=False,
                            weight_init=init.Normal(0.02, 0.0))]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_size, affine=False))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout(p=dropout))
        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x)


class UNetUp(nn.Cell):
    """U-Net Up"""

    def __init__(self, in_size, out_size, dropout=0.0):
        super().__init__(UNetUp)
        layers = [
            nn.Conv2dTranspose(in_size, out_size, 4, 2,
                               pad_mode='pad', padding=1, has_bias=False,
                               weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(out_size, affine=False),
            nn.ReLU()
        ]
        if dropout:
            layers.append(nn.Dropout(p=dropout))

        self.model = nn.SequentialCell(*layers)

    def construct(self, x, skip_input):
        x = self.model(x)
        x = ops.cat((x, skip_input), 1)
        return x


class GeneratorUNet(nn.Cell):
    """Generator Network"""

    def __init__(self, input_shape):
        super().__init__(GeneratorUNet)
        channels, _, _ = input_shape
        self.down1 = UNetDown(channels, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256, dropout=0.5)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = UNetDown(512, 512, dropout=0.5, normalize=False)

        self.up1 = UNetUp(512, 512, dropout=0.5)
        self.up2 = UNetUp(1024, 512, dropout=0.5)
        self.up3 = UNetUp(1024, 256, dropout=0.5)
        self.up4 = UNetUp(512, 128)
        self.up5 = UNetUp(256, 64)

        self.final = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(128, channels, 4, stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, x):
        # U-Net generator with skip connections from encoder to decoder
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        u1 = self.up1(d6, d5)
        u2 = self.up2(u1, d4)
        u3 = self.up3(u2, d3)
        u4 = self.up4(u3, d2)
        u5 = self.up5(u4, d1)

        return self.final(u5)


##############################
#        Discriminator
##############################


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, input_shape):
        super().__init__(Discriminator)

        channels, height, width = input_shape
        # Calculate output of image discriminator (PatchGAN)
        self.output_shape = (1, height // 2 ** 3, width // 2 ** 3)

        def discriminator_block(in_filters, out_filters, normalization=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalization:
                layers.append(nn.InstanceNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *discriminator_block(channels, 64, normalization=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(256, 1, 4, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0))
        )

    def construct(self, img):
        # Concatenate image and condition image by channels to produce input
        return self.model(img)

implementations/discogan/models.py
mindspore-courses/MindSpore-GAN
"""DRAGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=1000, help="interval between image sampling")
opt = parser.parse_args()
print(opt)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.init_size = opt.img_size // 4
        self.l1 = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 128 * self.init_size ** 2)
        )

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, stride=1,
                      pad_mode='pad', padding=1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh(),
        )

    def construct(self, noise):
        out = self.l1(noise)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [
                nn.Conv2d(in_filters, out_filters, 3, 2,
                          pad_mode='pad', padding=1, has_bias=False,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0), affine=False))
            return block

        self.model = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4
        self.adv_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, 1),
            nn.Sigmoid()
        )

    def construct(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)

        return validity


# Loss function
adversarial_loss = nn.BCELoss()

# Loss weight for gradient penalty
lambda_gp = 10

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def compute_gradient_penalty(D, X):
    """Calculates the gradient penalty loss for DRAGAN"""
    # Random weight term for interpolation
    alpha = Tensor(np.random.random(size=X.shape))

    interpolates = alpha * X + ((1 - alpha) * (X + 0.5 * X.std() * ops.rand(X.shape)))
    interpolates = ops.Cast()(interpolates, mstype.float32)

    # Get gradient w.r.t. interpolates

    grad_fn = ops.grad(D, return_ids=True)
    gradients = grad_fn(interpolates)
    gradients = ops.get_grad(gradients, 0)

    _gradient_penalty = lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return _gradient_penalty


def g_forward(_imgs, _valid):
    """Generator forward function"""
    # Sample noise as generator input
    z = ops.randn((_imgs.shape[0], opt.latent_dim), dtype=mstype.float32)

    # Generate a batch of images
    _gen_imgs = generator(z)

    # Loss measures generator's ability to fool the discriminator
    _g_loss = adversarial_loss(discriminator(_gen_imgs), _valid)
    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_loss = adversarial_loss(discriminator(_real_imgs), _valid)
    fake_loss = adversarial_loss(discriminator(_gen_imgs), _fake)
    _d_loss = (real_loss + fake_loss) / 2

    # Calculate gradient penalty
    _gradient_penalty = compute_gradient_penalty(discriminator, _real_imgs)

    return _gradient_penalty, _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=True)

generator.set_train()
discriminator.set_train()

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        valid = ops.stop_gradient(ops.ones((imgs.shape[0], 1)))
        fake = ops.stop_gradient(ops.zeros((imgs.shape[0], 1)))

        # Configure input
        real_imgs = imgs

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(real_imgs, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (gradient_penalty, d_loss), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs), valid, fake)
        optimizer_D(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}]'
        )

    to_image(gen_imgs, os.path.join("images", F'{epoch}.png'))

implementations/dragan/dragan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/dragan/img_utils.py
mindspore-courses/MindSpore-GAN
"""Dataset Preprocess"""

import glob
import os

import numpy as np
from PIL import Image
from mindspore.dataset.transforms import Compose


class Edges2ShoesDataset():
    """Edges2Shoes Dataset"""

    def __init__(self, root, transforms_=None, mode='train'):
        self.transform = Compose(transforms_)

        self.files = sorted(glob.glob(os.path.join(root, mode) + '/*.*'))

    def __getitem__(self, index):
        img = Image.open(self.files[index % len(self.files)])
        w, h = img.size
        img_A = img.crop((0, 0, w / 2, h))
        img_B = img.crop((w / 2, 0, w, h))

        if np.random.random() < 0.5:
            img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], "RGB")
            img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], "RGB")

        img_A = self.transform(img_A)
        img_B = self.transform(img_B)

        return img_A[0], img_B[0]

    def __len__(self):
        return len(self.files)

implementations/dualgan/datasets.py
mindspore-courses/MindSpore-GAN
"""DualGAN Model"""

import argparse
import datetime
import itertools
import os
import sys
import time

import mindspore
import numpy as np
from mindspore import nn, Tensor
from mindspore import ops
from mindspore.dataset.vision import transforms, Inter
import mindspore.common.dtype as mstype

from datasets import Edges2ShoesDataset
from img_utils import to_image
from models import Generator, Discriminator

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=8, help="size of the batches")
parser.add_argument("--dataset_name", type=str, default="edges2shoes", help="name of the dataset")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_size", type=int, default=128, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--n_critic", type=int, default=5, help="number of training steps for discriminator per iter")
parser.add_argument("--sample_interval", type=int, default=200, help="interval betwen image samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
opt = parser.parse_args()
print(opt)

# Create sample and checkpoint directories
os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

# Losses
cycle_loss = nn.L1Loss()

# Loss weights
lambda_adv = 1
lambda_cycle = 10
lambda_gp = 10

# Initialize generator and discriminator
G_AB = Generator()
G_BA = Generator()
D_A = Discriminator()
D_B = Discriminator()

G_AB.update_parameters_name("G_AB")
G_BA.update_parameters_name("G_BA")
D_A.update_parameters_name("G_AB")
D_B.update_parameters_name("G_BA")

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G_AB_{opt.epoch}.ckpt', G_AB)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G_BA_{opt.epoch}.ckpt', G_BA)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_A_{opt.epoch}.ckpt', D_A)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D_B_{opt.epoch}.ckpt', D_B)

# Optimizers
optimizer_G = nn.optim.Adam(itertools.chain(G_AB.trainable_params(), G_BA.trainable_params()),
                            learning_rate=opt.lr,
                            beta1=opt.b1, beta2=opt.b2)
optimizer_D_A = nn.optim.Adam(D_A.trainable_params(), learning_rate=opt.lr,
                              beta1=opt.b1, beta2=opt.b2)

optimizer_D_B = nn.optim.Adam(D_B.trainable_params(), learning_rate=opt.lr,
                              beta1=opt.b1, beta2=opt.b2)

# Dataset loader
transforms_ = [
    transforms.Resize((opt.img_size, opt.img_size), Inter.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
]

dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        transforms_=transforms_
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(opt.batch_size)

val_dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        transforms_=transforms_,
        mode='val'
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(16)


def compute_gradient_penalty(D, real_samples, fake_samples):
    """Calculates the gradient penalty"""
    # Random weight term for interpolation between real and fake samples
    alpha = Tensor(np.random.random(size=(real_samples.shape[0], 1, 1, 1)),dtype=mstype.float32)
    # Get random interpolation between real and fake samples
    interpolates = alpha * real_samples + ((1 - alpha) * fake_samples)
    # Get gradient w.r.t. interpolates
    grad_fn = mindspore.grad(D)
    gradients = grad_fn(interpolates)
    # gradients = mindspore.get_grad(gradients, 0)
    gradients = gradients.view(gradients.shape[0], -1)
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty


def sample_image(batches):
    """Saves a generated sample from the validation set"""
    imgs = next(val_dataset.create_tuple_iterator())
    G_AB.set_train(False)
    G_BA.set_train(False)
    _real_A = imgs[0]
    _fake_B = G_AB(_real_A)
    AB = ops.cat((_real_A, _fake_B), -2)
    _real_B = imgs[1]
    _fake_A = G_BA(_real_B)
    BA = ops.cat((_real_B, _fake_A), -2)
    img_sample = ops.cat((AB, BA), 0)
    to_image(img_sample, os.path.join(f'images/{opt.dataset_name}', f'{batches}.png'))


def g_forward(_real_A, _real_B):
    """Generator warmup forward func"""
    # Translate images to opposite domain
    _fake_A = G_BA(_real_B)
    _fake_B = G_AB(_real_A)

    # Reconstruct images
    recov_A = G_BA(_fake_B)
    recov_B = G_AB(_fake_A)

    # Adversarial loss
    _G_adv = -ops.mean(D_A(_fake_A)) - ops.mean(D_B(_fake_B))
    # Cycle loss
    _G_cycle = cycle_loss(recov_A, _real_A) + cycle_loss(recov_B, _real_B)
    # Total loss
    _G_loss = lambda_adv * _G_adv + lambda_cycle * _G_cycle

    return _G_loss, _G_cycle, _G_adv


def d_a_forward(_real_A, _fake_A):
    """Discriminator A forward function"""

    # Compute gradient penalty for improved wasserstein training
    gp_A = compute_gradient_penalty(D_A, _real_A, _fake_A)
    # Adversarial loss
    _D_A_loss = -ops.mean(D_A(_real_A)) + ops.mean(D_A(_fake_A)) + lambda_gp * gp_A

    return _D_A_loss


def d_b_forward(_real_B, _fake_B):
    """Discriminator B forward function"""

    # Compute gradient penalty for improved wasserstein training
    gp_B = compute_gradient_penalty(D_B, _real_B, _fake_B)
    # Adversarial loss
    _D_B_loss = -ops.mean(D_B(_real_B)) + ops.mean(D_B(_fake_B)) + lambda_gp * gp_B

    return _D_B_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d_a = ops.value_and_grad(d_a_forward, None, optimizer_D_A.parameters, has_aux=False)
grad_d_b = ops.value_and_grad(d_b_forward, None, optimizer_D_B.parameters, has_aux=False)

# ----------
#  Training
# ----------

batches_done = 0
prev_time = time.time()
for epoch in range(opt.n_epochs):
    # Model inputs
    for i, (real_A, real_B) in enumerate(dataset.create_tuple_iterator()):

        # ----------------------
        #  Train Discriminators
        # ----------------------

        D_A.set_train()
        D_B.set_train()

        # Generate a batch of images
        fake_A = ops.stop_gradient(G_BA(real_B))
        fake_B = ops.stop_gradient(G_AB(real_A))

        # ----------
        # Domain A
        # ----------

        (D_A_loss), d_a_grads = grad_d_a(real_A, fake_A)
        optimizer_D_A(d_a_grads)

        # ----------
        # Domain B
        # ----------

        (D_B_loss), d_b_grads = grad_d_b(real_B, fake_B)
        optimizer_D_B(d_b_grads)

        # Total loss
        D_loss = D_A_loss + D_B_loss

        if i % opt.n_critic == 0:
            # ------------------
            #  Train Generators
            # ------------------

            G_AB.set_train()
            G_BA.set_train()

            (G_loss, G_cycle, G_adv), g_grads = grad_g(real_A, real_B)
            optimizer_G(g_grads)

            # --------------
            #  Log Progress
            # --------------

            # Determine approximate time left
            batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
            time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time) / opt.n_critic)
            prev_time = time.time()

            sys.stdout.write(
                f'\r[Epoch {epoch}/{opt.n_epochs}] '
                f'[Batch {i}/{dataset.get_dataset_size()}] '
                f'[D loss: {D_loss.asnumpy().item():.4f}] '
                f'[G loss: {G_adv.asnumpy().item():.4f}, '
                f'cycle:{G_cycle.asnumpy().item():.4f}] '
                f'ETA: {time_left}'
            )

        if batches_done % opt.sample_interval == 0:
            sample_image(batches_done)

        batches_done += 1

        if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
            # Save model checkpoints
            mindspore.save_checkpoint(G_AB, f'saved_models/{opt.dataset_name}/G_AB_{epoch}.ckpt')
            mindspore.save_checkpoint(G_BA, f'saved_models/{opt.dataset_name}/G_BA_{epoch}.ckpt')
            mindspore.save_checkpoint(D_A, f'saved_models/{opt.dataset_name}/D_A_{epoch}.ckpt')
            mindspore.save_checkpoint(D_B, f'saved_models/{opt.dataset_name}/D_B_{epoch}.ckpt')

implementations/dualgan/dualgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/dualgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""DualGAN Networks"""

import mindspore.common.initializer as init
from mindspore import nn, ops


##############################
#           U-NET
##############################


class UNetDown(nn.Cell):
    """U-Net Down"""

    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
        super().__init__(UNetDown)
        layers = [nn.Conv2d(in_size, out_size, 4, stride=2, pad_mode='pad', padding=1, has_bias=False,
                            weight_init=init.Normal(0.02, 0.0))]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_size, affine=True))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout(p=dropout))
        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x)


class UNetUp(nn.Cell):
    """U-Net Up"""

    def __init__(self, in_size, out_size, dropout=0.0):
        super().__init__(UNetUp)
        layers = [
            nn.Conv2dTranspose(in_size, out_size, 4, 2,
                               pad_mode='pad', padding=1, has_bias=False,
                               weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(out_size, affine=True),
            nn.ReLU()
        ]
        if dropout:
            layers.append(nn.Dropout(p=dropout))

        self.model = nn.SequentialCell(*layers)

    def construct(self, x, skip_input):
        x = self.model(x)
        x = ops.cat((x, skip_input), 1)
        return x


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self, channels=3):
        super().__init__(Generator)

        self.down1 = UNetDown(channels, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = UNetDown(512, 512, dropout=0.5)
        self.down7 = UNetDown(512, 512, dropout=0.5, normalize=False)

        self.up1 = UNetUp(512, 512, dropout=0.5)
        self.up2 = UNetUp(1024, 512, dropout=0.5)
        self.up3 = UNetUp(1024, 512, dropout=0.5)
        self.up4 = UNetUp(1024, 256)
        self.up5 = UNetUp(512, 128)
        self.up6 = UNetUp(256, 64)

        self.final = nn.SequentialCell(
            nn.Conv2dTranspose(128, channels, 4, stride=2, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, x):
        # Propogate noise through fc layer and reshape to img shape
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        d7 = self.down7(d6)
        u1 = self.up1(d7, d6)
        u2 = self.up2(u1, d5)
        u3 = self.up3(u2, d4)
        u4 = self.up4(u3, d3)
        u5 = self.up5(u4, d2)
        u6 = self.up6(u5, d1)

        return self.final(u6)


##############################
#        Discriminator
##############################


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, in_channels=3):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, normalize=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalize:
                layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *discriminator_block(in_channels, 64, normalize=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(256, 1, 4, pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0))
        )

    def construct(self, img):
        # Concatenate image and condition image by channels to produce input
        return self.model(img)

implementations/dualgan/models.py
mindspore-courses/MindSpore-GAN
"""EBGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
from mindspore import nn
from mindspore import ops
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=62, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="number of image channels")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.init_size = opt.img_size // 4
        self.l1 = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 128 * self.init_size ** 2)
        )

        self.conv_blocks = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh(),
        )

    def construct(self, noise):
        out = self.l1(noise)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        # Upsampling
        self.down = nn.SequentialCell(
            nn.Conv2d(opt.channels, 64, 3,
                      stride=2, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.ReLU()
        )
        # Fully-connected layers
        self.down_size = opt.img_size // 2
        down_dim = 64 * (opt.img_size // 2) ** 2

        self.embedding = nn.Dense(down_dim, 32)

        self.fc = nn.SequentialCell(
            nn.BatchNorm1d(32, 0.8, affine=False),
            nn.ReLU(),
            nn.Dense(32, down_dim),
            nn.BatchNorm1d(down_dim, affine=False),
            nn.ReLU(),
        )
        # Upsampling
        self.up = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(64, opt.channels, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0))
        )

    def construct(self, img):
        out = self.down(img)
        embedding = self.embedding(out.view(out.shape[0], -1))
        out = self.fc(embedding)
        out = self.up(out.view(out.shape[0], 64, self.down_size, self.down_size))
        return out, embedding


# Reconstruction loss of AE
pixelwise_loss = nn.MSELoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def pullaway_loss(embeddings):
    """Pull away loss"""
    norm = ops.sqrt(ops.sum(embeddings ** 2, -1, keepdim=True))
    normalized_emb = embeddings / norm
    similarity = ops.matmul(normalized_emb, normalized_emb.transpose(1, 0))
    batch_size = embeddings.shape[0]
    loss_pt = (ops.sum(similarity) - batch_size) / (batch_size * (batch_size - 1))
    return loss_pt


def g_forward(_imgs):
    """Generator forward function"""
    # Sample noise as generator input
    z = ops.randn((_imgs.shape[0], opt.latent_dim), dtype=mstype.float32)

    # Generate a batch of images
    _gen_imgs = generator(z)
    recon_imgs, img_embeddings = discriminator(_gen_imgs)

    # Loss measures generator's ability to fool the discriminator
    _g_loss = pixelwise_loss(recon_imgs, _gen_imgs) + lambda_pt * pullaway_loss(img_embeddings)

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_recon, _ = discriminator(_real_imgs)
    fake_recon, _ = discriminator(_gen_imgs)

    d_loss_real = pixelwise_loss(real_recon, _real_imgs)
    d_loss_fake = pixelwise_loss(fake_recon, _gen_imgs)

    _d_loss = d_loss_real
    if (margin - d_loss_fake).asnumpy().item() > 0:
        _d_loss += margin - d_loss_fake

    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

generator.set_train()
discriminator.set_train()

# ----------
#  Training
# ----------

# BEGAN hyper parameters
lambda_pt = 0.1
margin = max(1, opt.batch_size / 64.0)

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        # Configure input
        real_imgs = imgs

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(real_imgs)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (d_loss), d_grads = grad_d(real_imgs, gen_imgs)
        optimizer_D(d_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/ebgan/ebgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/ebgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""
Super-resolution of CelebA using Generative Adversarial Networks.
The dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0
(if not available there see if options are listed at http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)
Instrustion on running the script:
1. Download the dataset from the provided link
2. Save the folder 'img_align_celeba' to '../../data/'
4. Run the sript using command 'python3 esrgan.py'
"""

import argparse
import os
import sys

import mindspore
import numpy as np
from mindspore import nn
from mindspore import ops
from mindspore.dataset import CelebADataset
from mindspore.dataset.vision import transforms, Inter

from img_utils import to_image
from models import GeneratorRRDB, Discriminator, FeatureExtractor

os.makedirs("images/training", exist_ok=True)
os.makedirs("saved_models", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=4, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.9, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--hr_height", type=int, default=256, help="high res. image height")
parser.add_argument("--hr_width", type=int, default=256, help="high res. image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=100, help="interval between saving image samples")
parser.add_argument("--checkpoint_interval", type=int, default=5000, help="batch interval between model checkpoints")
parser.add_argument("--residual_blocks", type=int, default=23, help="number of residual blocks in the generator")
parser.add_argument("--warmup_batches", type=int, default=500, help="number of batches with pixel-wise loss only")
parser.add_argument("--lambda_adv", type=float, default=5e-3, help="adversarial loss weight")
parser.add_argument("--lambda_pixel", type=float, default=1e-2, help="pixel-wise loss weight")
opt = parser.parse_args()
print(opt)

hr_shape = (opt.hr_height, opt.hr_width)

# Initialize generator and discriminator
generator = GeneratorRRDB(opt.channels, filters=64, num_res_blocks=opt.residual_blocks)
discriminator = Discriminator(input_shape=(opt.channels, *hr_shape))
feature_extractor = FeatureExtractor()

# Set feature extractor to inference mode
feature_extractor.set_train(False)

criterion_GAN = nn.BCEWithLogitsLoss()
criterion_content = nn.L1Loss()
criterion_pixel = nn.L1Loss()

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/generator_{opt.epoch}.ckpt', generator)
    mindspore.load_checkpoint(f'saved_models/discriminator{opt.epoch}.ckpt', discriminator)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), opt.lr, opt.b1, opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), opt.lr, opt.b1, opt.b2)

mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])

def denormalize(tensors):
    """ Denormalizes image tensors using mean and std """
    for c in range(3):
        tensors[:, c].mul(std[c]).add(mean[c])
    return ops.clamp(tensors, 0, 255)

def preprocess(_imgs):
    """Dataset preprocess func"""
    _imgs = transforms.Resize((opt.hr_height, opt.hr_width), Inter.BICUBIC)(_imgs)
    _imgs_lr = transforms.Resize((opt.hr_height // 4, opt.hr_width // 4), Inter.BICUBIC)(_imgs)
    _imgs = transforms.ToTensor()(_imgs)
    _imgs = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], is_hwc=False)(_imgs)
    _imgs_lr = transforms.ToTensor()(_imgs_lr)
    _imgs_lr = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], is_hwc=False)(_imgs_lr)

    return _imgs, _imgs_lr


dataset = CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True
).map(operations=[preprocess], input_columns=["image"], output_columns=["image_hr", "image_lr"]).batch(opt.batch_size)


def g_forward(_imgs_hr, _imgs_lr, _valid):
    """Generator forward function"""
    # Generate a high resolution image from low resolution input
    _gen_hr = generator(_imgs_lr)
    # Measure pixel-wise loss against ground truth
    _pixel_loss = criterion_pixel(_gen_hr, _imgs_hr)

    # Extract validity predictions from discriminator
    pred_real = ops.stop_gradient(discriminator(_imgs_hr))
    pred_fake = discriminator(_gen_hr)

    # Adversarial loss (relativistic average GAN)
    _GAN_loss = criterion_GAN(pred_fake - pred_real.mean(0, keep_dims=True), _valid)

    # Content loss
    gen_features = feature_extractor(_gen_hr)
    real_features = ops.stop_gradient(feature_extractor(_imgs_hr))
    _content_loss = criterion_content(gen_features, real_features)

    # Total generator loss
    _g_loss = _content_loss + opt.lambda_adv * _GAN_loss + opt.lambda_pixel * loss_pixel

    return _g_loss, _content_loss, _GAN_loss, _pixel_loss, _gen_hr


def g_forward_warmup(_imgs_hr, _imgs_lr):
    """Generator warmup forward func"""
    # Generate a high resolution image from low resolution input
    _gen_hr = generator(_imgs_lr)
    # Measure pixel-wise loss against ground truth
    _loss_pixel = criterion_pixel(_gen_hr, _imgs_hr)
    return _loss_pixel


def d_forward(_imgs_hr, _gen_hr, _valid, _fake):
    """Discriminator forward function"""
    pred_real = discriminator(_imgs_hr)
    pred_fake = discriminator(_gen_hr)

    # Adversarial loss for real and fake images (relativistic average GAN)
    loss_real = criterion_GAN(pred_real - pred_fake.mean(0, keep_dims=True), _valid)
    loss_fake = criterion_GAN(pred_fake - pred_real.mean(0, keep_dims=True), _fake)

    # Total loss
    _d_loss = (loss_real + loss_fake) / 2
    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)
grad_g_warmup = ops.value_and_grad(g_forward_warmup, None, optimizer_G.parameters, has_aux=False)
# ----------
#  Training
# ----------

saved_samples = {}

for epoch in range(opt.n_epochs):
    # Model inputs
    for i, batch in enumerate(dataset.create_dict_iterator()):

        batches_done = epoch * dataset.get_dataset_size() + i

        imgs_hr = batch["image_hr"]
        imgs_lr = batch["image_lr"]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs_lr.shape[0], *discriminator.output_shape)))
        fake = ops.stop_gradient(ops.zeros((imgs_lr.shape[0], *discriminator.output_shape)))

        # ------------------
        #  Train Generators
        # ------------------

        if batches_done < opt.warmup_batches:
            # Warm-up (pixel-wise loss only)
            (loss_pixel), g_grads = grad_g_warmup(imgs_hr, imgs_lr)
            optimizer_G(g_grads)
            print(
                f'[Epoch {epoch}/{opt.n_epochs}] '
                f'[Batch {i}/{dataset.get_dataset_size()}] '
                f'[G pixel: {loss_pixel.asnumpy().item()}]'
            )
            continue

        (g_loss, content_loss, GAN_loss, pixel_loss, gen_hr), g_grads = grad_g(imgs_hr, imgs_lr, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss), d_grads = grad_d(imgs_hr, ops.stop_gradient(gen_hr),
                                   valid, fake)
        optimizer_D(d_grads)

        # --------------
        #  Log Progress
        # --------------

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G loss: {g_loss.asnumpy().item():.4f}, '
            f'content: {content_loss.asnumpy().item():.4f}, '
            f'adv: {GAN_loss.asnumpy().item():.4f}, '
            f'pixel:{pixel_loss.asnumpy().item():.4f}]'
        )

        # If at sample interval save image
        if batches_done % opt.sample_interval == 0:
            # Save image grid with upsampled inputs and SRGAN outputs
            imgs_lr = ops.interpolate(imgs_lr, scale_factor=4.0, recompute_scale_factor=True)
            img_grid = denormalize(ops.cat((imgs_lr, gen_hr), -1))
            to_image(img_grid, f'images/training/{batches_done}.png')

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        mindspore.save_checkpoint(generator, f'saved_models/{opt.dataset_name}/generator_{epoch}.ckpt')
        mindspore.save_checkpoint(discriminator, f'saved_models/{opt.dataset_name}/discriminator_{epoch}.ckpt')

implementations/esrgan/esrgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 1,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/esrgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""The Models of ESRGAN"""

from mindcv.models import vgg19
from mindspore import nn
from mindspore import ops


class FeatureExtractor(nn.Cell):
    """VGG Feature Extractor"""

    def __init__(self):
        super().__init__(FeatureExtractor)
        vgg19_model = vgg19(pretrained=True)
        self.vgg19_54 = nn.SequentialCell(
            *list(vgg19_model.features.cell_list[:35])
        )

    def construct(self, img):
        return self.vgg19_54(img)


class DenseResidualBlock(nn.Cell):
    """
    The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)
    """

    def __init__(self, filters, res_scale=0.2):
        super().__init__(DenseResidualBlock)
        self.res_scale = res_scale

        def block(in_features, non_linearity=True):
            layers = [nn.Conv2d(in_features, filters, 3, 1,
                                'pad', 1, has_bias=True)]
            if non_linearity:
                layers += [nn.LeakyReLU()]
            return nn.SequentialCell(*layers)

        self.b1 = block(in_features=1 * filters)
        self.b2 = block(in_features=2 * filters)
        self.b3 = block(in_features=3 * filters)
        self.b4 = block(in_features=4 * filters)
        self.b5 = block(in_features=5 * filters, non_linearity=False)
        self.blocks = [self.b1, self.b2, self.b3, self.b4, self.b5]

    def construct(self, x):
        inputs = x
        for block in self.blocks:
            out = block(inputs)
            inputs = ops.cat([inputs, out], 1)
        return out.mul(self.res_scale) + x


class ResidualInResidualDenseBlock(nn.Cell):
    """Residual In Residual Block"""
    def __init__(self, filters, res_scale=0.2):
        super().__init__(ResidualInResidualDenseBlock)
        self.res_scale = res_scale
        self.dense_blocks = nn.SequentialCell(
            DenseResidualBlock(filters),
            DenseResidualBlock(filters),
            DenseResidualBlock(filters)
        )

    def construct(self, x):
        return self.dense_blocks(x).mul(self.res_scale) + x


class GeneratorRRDB(nn.Cell):
    """ResNet Generator"""

    def __init__(self, channels, filters=64, num_res_blocks=16, num_upsample=2):
        super().__init__(GeneratorRRDB)

        # First layer
        self.conv1 = nn.Conv2d(channels, filters, 3, 1,
                               'pad', 1)

        # Residual blocks
        self.res_blocks = nn.SequentialCell(
            *[ResidualInResidualDenseBlock(filters) for _ in range(num_res_blocks)]
        )

        # Second conv layer post residual blocks
        self.conv2 = nn.Conv2d(filters, filters, 3, 1,
                               'pad', 1)

        # Upsampling layers
        upsample_layers = []
        for _ in range(num_upsample):
            upsample_layers += [
                nn.Conv2d(filters, filters * 4, 3, 1,
                          'pad', 1),
                nn.LeakyReLU(),
                nn.PixelShuffle(upscale_factor=2),
            ]
        self.upsampling = nn.SequentialCell(*upsample_layers)

        # Final output layer
        self.conv3 = nn.SequentialCell(
            nn.Conv2d(filters, filters, 3, 1, 'pad', 1),
            nn.LeakyReLU(),
            nn.Conv2d(filters, channels, 3, 1, 'pad', 1),
        )

    def construct(self, x):
        out1 = self.conv1(x)
        out = self.res_blocks(out1)
        out2 = self.conv2(out)
        out = ops.add(out1, out2)
        out = self.upsampling(out)
        out = self.conv3(out)
        return out


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, input_shape):
        super().__init__(Discriminator)

        self.input_shape = input_shape
        in_channels, in_height, in_width = self.input_shape
        patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)
        self.output_shape = (1, patch_h, patch_w)

        def discriminator_block(in_filters, out_filters, first_block=False):
            layers = []
            layers.append(nn.Conv2d(in_filters, out_filters, 3,
                                    1, 'pad', 1))
            if not first_block:
                layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            layers.append(nn.Conv2d(out_filters, out_filters, 3,
                                    2, 'pad', 1))
            layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        layers = []
        in_filters = in_channels
        for i, out_filters in enumerate([64, 128, 256, 512]):
            layers.extend(discriminator_block(in_filters, out_filters, first_block=i == 0))
            in_filters = out_filters

        layers.append(nn.Conv2d(out_filters, 1, 3,
                                1, 'pad', 1))

        self.model = nn.SequentialCell(*layers)

    def construct(self, img):
        return self.model(img)

implementations/esrgan/models.py
mindspore-courses/MindSpore-GAN
"""Test ESRGAN"""
import argparse
import os

import mindspore
import numpy as np
from mindspore import Tensor, ops
from mindspore.dataset.transforms import Compose
from mindspore.dataset.vision import transforms

from PIL import Image

from models import GeneratorRRDB
from img_utils import to_image

parser = argparse.ArgumentParser()
parser.add_argument("--image_path", type=str, required=True, help="Path to image")
parser.add_argument("--checkpoint_model", type=str, required=True, help="Path to checkpoint model")
parser.add_argument("--channels", type=int, default=3, help="Number of image channels")
parser.add_argument("--residual_blocks", type=int, default=23, help="Number of residual blocks in G")
opt = parser.parse_args()
print(opt)

mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])


def denormalize(tensors):
    """ Denormalizes image tensors using mean and std """
    for c in range(3):
        tensors[:, c].mul(std[c]).add(mean[c])
    return ops.clamp(tensors, 0, 255)


os.makedirs("images/outputs", exist_ok=True)
# Define model and load model checkpoint
generator = GeneratorRRDB(opt.channels, filters=64, num_res_blocks=opt.residual_blocks)
mindspore.load_checkpoint(opt.checkpoint_model, generator)
generator.set_train(False)

transform = Compose(
    [transforms.ToTensor(), transforms.Normalize(mean, std, is_hwc=False)]
)

# Prepare input
image_tensor = Tensor(transform(Image.open(opt.image_path))).unsqueeze(0)

sr_image = denormalize(ops.stop_gradient(generator(image_tensor)))

fn = opt.image_path.split("/")[-1]
to_image(sr_image, f'images/outputs/sr-{fn}')

implementations/esrgan/test_on_image.py
mindspore-courses/MindSpore-GAN
"""GAN Model"""

import argparse
import gzip
import math
import os
import shutil
import urllib.request

import numpy as np
import mindspore
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common.initializer import HeUniform
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size",
                    type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, x):
        return self.model(x)


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)
        self.fc1 = nn.Dense(int(np.prod(img_shape)), 512, weight_init=HeUniform(math.sqrt(5)))
        self.relu = nn.LeakyReLU(0.2)
        self.fc2 = nn.Dense(512, 256, weight_init=HeUniform(math.sqrt(5)))
        self.fc3 = nn.Dense(256, 1, weight_init=HeUniform(math.sqrt(5)))
        self.sigmoid = nn.Sigmoid()

    def construct(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.sigmoid(x)
        return x


generator = Generator()
discriminator = Discriminator()

# 损失函数与优化器
criterion = nn.BCELoss()
D_Optim = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
G_Optim = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_z):
    """Generator forward function"""
    img = generator(_z)
    img = img.view(img.shape[0], *img_shape)
    img_flat = img.view(img.shape[0], -1)
    dis = discriminator(img_flat)
    _g_loss = criterion(dis, valid)
    return _g_loss, img


# 判别器正向传播
def d_forward(img, gen_img, _valid, _fake):
    """Discriminator forward function"""
    img_flat = img.view(img.shape[0], -1)
    gen_flat = gen_img.view(img.shape[0], -1)
    # _real_score = discriminator(img)
    # _fake_score = discriminator(img_flat)
    real_loss = criterion(discriminator(img_flat), _valid)
    fake_loss = criterion(discriminator(gen_flat), _fake)
    _d_loss = real_loss + fake_loss
    return _d_loss


transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, G_Optim.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, D_Optim.parameters, has_aux=False)

# 训练
for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        generator.set_train()
        discriminator.set_train()
        valid = ops.ones((imgs.shape[0], 1))
        fake = ops.zeros((imgs.shape[0], 1))
        ops.stop_gradient(valid)
        ops.stop_gradient(fake)

        real_imgs = Tensor(imgs)

        z = ops.randn((imgs.shape[0], opt.latent_dim))

        (g_loss, gen_images), g_grads = grad_g(z)
        G_Optim(g_grads)

        (d_loss), d_grads = grad_d(imgs, gen_images, valid, fake)
        D_Optim(d_grads)

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item()}] [G loss: {g_loss.asnumpy().item()}]'
        )

        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            to_image(gen_images[:25], os.path.join("images", F'{batches_done}.png'))

implementations/gan/gan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = False,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/gan/img_utils.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 10,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/infogan/img_utils.py
mindspore-courses/MindSpore-GAN
"""InfoGAN Model"""

import argparse
import gzip
import itertools
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images/static/", exist_ok=True)
os.makedirs("images/varying_c1/", exist_ok=True)
os.makedirs("images/varying_c2/", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=62, help="dimensionality of the latent space")
parser.add_argument("--code_dim", type=int, default=2, help="latent code")
parser.add_argument("--n_classes", type=int, default=10, help="number of classes for dataset")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)


def to_categorical(y, num_columns):
    """Returns one-hot encoded Variable"""
    y_cat = np.zeros((y.shape[0], num_columns))
    y_cat[range(y.shape[0]), y] = 1.0

    return Tensor(y_cat, dtype=mstype.float32)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)
        input_dim = opt.latent_dim + opt.n_classes + opt.code_dim

        self.init_size = opt.img_size // 4  # Initial size before upsampling
        self.l1 = nn.SequentialCell(
            nn.Dense(input_dim, 128 * self.init_size ** 2)
        )

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh(),
        )

    def construct(self, noise, _labels, code):
        gen_input = ops.cat((noise, _labels, code), -1)
        out = self.l1(gen_input)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""
    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            """Returns layers of each discriminator block"""
            block = [
                nn.Conv2d(in_filters, out_filters, 3,
                          stride=2, pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0), affine=False))
            return block

        self.conv_blocks = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4

        # Output layers
        self.adv_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, 1)
        )
        self.aux_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, opt.n_classes),
            nn.Softmax())
        self.latent_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, opt.code_dim)
        )

    def construct(self, img):
        out = self.conv_blocks(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        label = self.aux_layer(out)
        latent_code = self.latent_layer(out)

        return validity, label, latent_code


# Loss functions
adversarial_loss = nn.MSELoss()
categorical_loss = nn.CrossEntropyLoss()
continuous_loss = nn.MSELoss()

lambda_cat = 1
lambda_con = 0.1

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_info = nn.optim.Adam(
    itertools.chain(generator.trainable_params(), discriminator.trainable_params()),
    learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2
)

# Static generator inputs for sampling
static_z = ops.zeros((opt.n_classes ** 2, opt.latent_dim), dtype=mstype.float32)
static_label = to_categorical(
    np.array([num for _ in range(opt.n_classes) for num in range(opt.n_classes)]), num_columns=opt.n_classes
)
static_code = ops.zeros((opt.n_classes ** 2, opt.code_dim), dtype=mstype.float32)


def sample_image(_n_row, batches):
    """Saves a grid of generated digits ranging from 0 to n_classes"""
    # Static sample
    z = ops.randn((_n_row ** 2, opt.latent_dim), dtype=mstype.float32)
    static_sample = generator(z, static_label, static_code)
    to_image(static_sample, f'images/static/{batches}.png')

    # Get varied c1 and c2
    zeros = np.zeros((_n_row ** 2, 1))
    c_varied = np.repeat(np.linspace(-1, 1, _n_row)[:, np.newaxis], _n_row, 0)
    c1 = Tensor(np.concatenate((c_varied, zeros), -1), dtype=mstype.float32)
    c2 = Tensor(np.concatenate((zeros, c_varied), -1), dtype=mstype.float32)
    sample1 = generator(static_z, static_label, c1)
    sample2 = generator(static_z, static_label, c2)
    to_image(sample1, f'images/varying_c1/{batches}.png')
    to_image(sample2, f'images/varying_c2/{batches}.png')


def g_forward(_imgs, _valid):
    """Generator forward function"""
    # Sample noise and labels as generator input
    z = ops.randn((batch_size, opt.latent_dim), dtype=mstype.float32)
    label_input = to_categorical(np.random.randint(0, opt.n_classes, batch_size), num_columns=opt.n_classes)
    code_input = ops.uniform((batch_size, opt.code_dim), Tensor(-1.0), Tensor(1.0))

    # Generate a batch of images
    _gen_imgs = generator(z, label_input, code_input)

    # Loss measures generator's ability to fool the discriminator
    validity, _, _ = discriminator(_gen_imgs)
    _g_loss = adversarial_loss(validity, _valid)
    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""
    # Loss for real images
    real_pred, _, _ = discriminator(_real_imgs)
    d_real_loss = adversarial_loss(real_pred, _valid)

    # Loss for fake images
    fake_pred, _, _ = discriminator(_gen_imgs)
    d_fake_loss = adversarial_loss(fake_pred, _fake)

    # Total discriminator loss
    _d_loss = (d_real_loss + d_fake_loss) / 2

    return _d_loss


def info_forward(_batch_size):
    """Information"""
    # Sample labels
    sampled_labels = np.random.randint(0, opt.n_classes, _batch_size)

    # Ground truth labels
    gt_labels = ops.stop_gradient(Tensor(sampled_labels, dtype=mstype.int32))

    # Sample noise, labels and code as generator input
    z = ops.randn((_batch_size, opt.latent_dim), dtype=mstype.float32)
    label_input = to_categorical(sampled_labels, num_columns=opt.n_classes)
    code_input = ops.uniform((_batch_size, opt.code_dim), Tensor(-1.0), Tensor(1.0))

    _gen_imgs = generator(z, label_input, code_input)
    _, pred_label, pred_code = discriminator(_gen_imgs)

    _info_loss = lambda_cat * categorical_loss(pred_label, gt_labels) + lambda_con * continuous_loss(
        pred_code, code_input
    )

    return _info_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)
grad_info = ops.value_and_grad(info_forward, None, optimizer_info.parameters, has_aux=False)

generator.set_train()
discriminator.set_train()

# ----------
#  Training
# ----------

for epoch in range(opt.n_epochs):
    for i, (imgs, labels) in enumerate(dataset.create_tuple_iterator()):
        batch_size = imgs.shape[0]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((batch_size, 1)))
        fake = ops.stop_gradient(ops.zeros((batch_size, 1)))

        # Configure input
        real_imgs = imgs
        labels = to_categorical(labels.asnumpy(), num_columns=opt.n_classes)

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(real_imgs, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (d_loss), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs), valid, fake)
        optimizer_D(d_grads)

        # ------------------
        # Information Loss
        # ------------------

        (info_loss), info_grads = grad_info(batch_size)
        optimizer_info(info_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}] '
            f'[info loss: {info_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            sample_image(_n_row=10, batches=batches_done)

implementations/infogan/infogan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/lsgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""LSGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
from mindspore import nn
from mindspore import ops
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=1000, help="number of image channels")
opt = parser.parse_args()
print(opt)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.init_size = opt.img_size // 4
        self.l1 = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 128 * self.init_size ** 2)
        )

        self.conv_blocks = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3, stride=1, padding=1,
                      pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3, stride=1, padding=1,
                      pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, stride=1, padding=1,
                      pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, z):
        out = self.l1(z)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [
                nn.Conv2d(in_filters, out_filters, 3, 2, padding=1,
                          pad_mode='pad',
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0)))
            return block

        self.model = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4
        self.adv_layer = nn.Dense(128 * ds_size ** 2, 1)

    def construct(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)

        return validity


# !!! Minimizes MSE instead of BCE
adversarial_loss = nn.MSELoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_imgs, _valid):
    """Generator forward function"""
    # Sample noise as generator input
    z = ops.randn((_imgs.shape[0], opt.latent_dim), dtype=mstype.float32)

    # Generate a batch of images
    _gen_imgs = generator(z)

    # Loss measures generator's ability to fool the discriminator
    _g_loss = adversarial_loss(discriminator(_gen_imgs), _valid)

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_loss = adversarial_loss(discriminator(_real_imgs), _valid)
    fake_loss = adversarial_loss(discriminator(_gen_imgs), _fake)
    _d_loss = 0.5 * (real_loss + fake_loss)

    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

generator.set_train()
discriminator.set_train()

# ----------
#  Training
# ----------

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs.shape[0], 1)))
        fake = ops.stop_gradient(ops.zeros((imgs.shape[0], 1)))

        # Configure input
        real_imgs = imgs

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(real_imgs, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (d_loss), d_grads = grad_d(real_imgs, gen_imgs, valid, fake)
        optimizer_D(d_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] [G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/lsgan/lsgan.py
mindspore-courses/MindSpore-GAN
"""Dataset Preprocess"""

import glob
import os

import numpy as np
from PIL import Image
from mindspore.dataset.transforms import Compose


class Edges2ShoesDataset():
    """Edges2Shoes Dataset"""

    def __init__(self, root, transforms_=None, mode='train'):
        self.transform = Compose(transforms_)

        self.files = sorted(glob.glob(os.path.join(root, mode) + '/*.*'))

    def __getitem__(self, index):
        img = Image.open(self.files[index % len(self.files)])
        w, h = img.size
        img_A = img.crop((0, 0, w / 2, h))
        img_B = img.crop((w / 2, 0, w, h))

        if np.random.random() < 0.5:
            img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], "RGB")
            img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], "RGB")

        img_A = self.transform(img_A)
        img_B = self.transform(img_B)

        return img_A[0], img_B[0]

    def __len__(self):
        return len(self.files)

implementations/munit/datasets.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/munit/img_utils.py
mindspore-courses/MindSpore-GAN
"""MUNIT Models"""

from mindspore import nn, ops, Parameter, Tensor
import mindspore.common.initializer as init
from mindspore.nn.learning_rate_schedule import LearningRateSchedule
import mindspore.common.dtype as mstype


class LambdaLR(LearningRateSchedule):
    """Learning rate decay"""

    def __init__(self, lr, n_epochs, offset, decay_start_epoch, step_per_epoch):
        super().__init__()
        assert (n_epochs - decay_start_epoch) > 0, "Decay must start before the training session ends!"
        self.lr = lr
        self.n_epochs = n_epochs
        self.offset = offset
        self.decay_start_epoch = decay_start_epoch
        self.cast = ops.Cast()
        self.step_per_epoch = step_per_epoch

    def construct(self, global_step):
        epoch = self.cast(global_step, mstype.float32) // self.step_per_epoch
        return self.lr * (1.0 - max(Tensor(0.0), epoch + self.offset - self.decay_start_epoch) / (
                self.n_epochs - self.decay_start_epoch))


#################################
#           Encoder
#################################


class Encoder(nn.Cell):
    """Encoder Network"""

    def __init__(self, in_channels=3, dim=64, n_residual=3, n_downsample=2, style_dim=8):
        super().__init__(Encoder)
        self.content_encoder = ContentEncoder(in_channels, dim, n_residual, n_downsample)
        self.style_encoder = StyleEncoder(in_channels, dim, n_downsample, style_dim)

    def construct(self, x):
        content_code = self.content_encoder(x)
        style_code = self.style_encoder(x)
        return content_code, style_code


#################################
#            Decoder
#################################


class Decoder(nn.Cell):
    """Decoder Network"""

    def __init__(self, out_channels=3, dim=64, n_residual=3, n_upsample=2, style_dim=8):
        super().__init__(Decoder)

        layers = []
        dim = dim * 2 ** n_upsample
        # Residual blocks
        for _ in range(n_residual):
            layers += [ResidualBlock(dim, norm="adain")]

        # Upsampling
        for _ in range(n_upsample):
            layers += [
                nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
                nn.Conv2d(dim, dim // 2, 5, stride=1, pad_mode='pad', padding=2,
                          weight_init=init.Normal(0.02, 0.0)),
                LayerNorm(dim // 2),
                nn.ReLU()
            ]
            dim = dim // 2

        # Output layer
        layers += [
            nn.ReflectionPad2d((3, 3, 3, 3)),
            nn.Conv2d(dim, out_channels, 7, pad_mode='valid',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        ]

        self.model = nn.SequentialCell(*layers)

        # Initiate mlp (predicts AdaIN parameters)
        num_adain_params = self.get_num_adain_params()
        self.mlp = MLP(style_dim, num_adain_params)

    def get_num_adain_params(self):
        """Return the number of AdaIN parameters needed by the model"""
        num_adain_params = 0
        for _, cell in self.cells_and_names():
            if cell.cls_name == "AdaptiveInstanceNorm2d":
                num_adain_params += 2 * cell.num_features
        return num_adain_params

    def assign_adain_params(self, adain_params):
        """Assign the adain_params to the AdaIN layers in model"""
        for _, cell in self.cells_and_names():
            if cell.cls_name == "AdaptiveInstanceNorm2d":
                # Extract mean and std predictions
                mean = adain_params[:, : cell.num_features]
                std = adain_params[:, cell.num_features: 2 * cell.num_features]
                # Update bias and weight
                cell.bias = mean.view(-1)
                cell.weight = std.view(-1)
                # Move pointer
                if adain_params.shape[1] > 2 * cell.num_features:
                    adain_params = adain_params[:, 2 * cell.num_features:]

    def construct(self, content_code, style_code):
        # Update AdaIN parameters by MLP prediction based off style code
        self.assign_adain_params(self.mlp(style_code))
        img = self.model(content_code)
        return img


#################################
#        Content Encoder
#################################

class ContentEncoder(nn.Cell):
    """Content Encoder"""

    def __init__(self, in_channels=3, dim=64, n_residual=3, n_downsample=2):
        super().__init__(ContentEncoder)

        # Initial convolution block
        layers = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(in_channels, dim, 7, pad_mode='valid',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(dim, affine=False),
            nn.ReLU()
        ]

        # Downsampling
        for _ in range(n_downsample):
            layers += [
                nn.Conv2d(dim, dim * 2, 4, stride=2, pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.InstanceNorm2d(dim * 2, affine=False),
                nn.ReLU()
            ]
            dim *= 2

        # Residual blocks
        for _ in range(n_residual):
            layers += [ResidualBlock(dim, norm="in")]

        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x)


#################################
#        Style Encoder
#################################


class StyleEncoder(nn.Cell):
    """Style Encoder"""
    def __init__(self, in_channels=3, dim=64, n_downsample=2, style_dim=8):
        super().__init__(StyleEncoder)

        # Initial conv block
        layers = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(in_channels, dim, 7, pad_mode='valid',
                      weight_init=init.Normal(0.2, 0.0)),
            nn.ReLU()
        ]

        # Downsampling
        for _ in range(2):
            layers += [
                nn.Conv2d(dim, dim * 2, 4, stride=2, pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.2, 0.0)),
                nn.ReLU()
            ]
            dim *= 2

        # Downsampling with constant depth
        for _ in range(n_downsample - 2):
            layers += [
                nn.Conv2d(dim, dim, 4, stride=2, padding=1, pad_mode='pad',
                          weight_init=init.Normal(0.2, 0.0)),
                nn.ReLU()
            ]

        # Average pool and output layer
        layers += [
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(dim, style_dim, 1, 1,
                      'pad', 0,
                      weight_init=init.Normal(0.2, 0.0))]

        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        # x = nn.ReflectionPad2d(3)(x)
        # x = nn.Conv2d(3, 64, 7, pad_mode='valid',
        #           weight_init=init.Normal(0.2, 0.0))(x)
        # x = nn.ReLU()(x)
        # x = nn.Conv2d(128, 128, 4, stride=2, padding=1, pad_mode='pad',
        #           weight_init=init.Normal(0.2, 0.0))(x)
        return self.model(x)


######################################
#   MLP (predicts AdaIn parameters)
######################################


class MLP(nn.Cell):
    """MLP model"""

    def __init__(self, input_dim, output_dim, dim=256, n_blk=3):
        super().__init__(MLP)
        layers = [nn.Dense(input_dim, dim), nn.ReLU()]
        for _ in range(n_blk - 2):
            layers += [nn.Dense(dim, dim), nn.ReLU()]
        layers += [nn.Dense(dim, output_dim)]
        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x.view(x.shape[0], -1))


##############################
#        Discriminator
##############################


class MultiDiscriminator(nn.Cell):
    """Multi-Discriminator"""
    def __init__(self, in_channels=3):
        super().__init__(MultiDiscriminator)

        def discriminator_block(in_filters, out_filters, normalize=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2,
                                pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.2, 0.0))]
            if normalize:
                layers.append(nn.InstanceNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        # Extracts three discriminator models
        self.models = nn.CellList()
        for _ in range(3):
            self.models.append(
                nn.SequentialCell(
                    *discriminator_block(in_channels, 64, normalize=False),
                    *discriminator_block(64, 128),
                    *discriminator_block(128, 256),
                    *discriminator_block(256, 512),
                    nn.Conv2d(512, 1, 3,
                              padding=1, pad_mode='pad',
                              weight_init=init.Normal(0.2, 0.0))
                )
            )

        self.downsample = nn.AvgPool2d(in_channels, stride=2, padding=[1, 1], pad_mode='pad', count_include_pad=False)

    def compute_loss(self, x, gt):
        """Computes the MSE between model output and scalar gt"""
        loss = sum([ops.mean((out - gt) ** 2) for out in self.construct(x)])
        return loss

    def construct(self, x):
        outputs = []
        for m in self.models:
            outputs.append(m(x))
            x = self.downsample(x)
        return outputs


##############################
#       Custom Blocks
##############################


class ResidualBlock(nn.Cell):
    """Residual blocks"""

    def __init__(self, features, norm="in"):
        super().__init__(ResidualBlock)

        norm_layer = AdaptiveInstanceNorm2d if norm == "adain" else nn.InstanceNorm2d

        self.block = nn.SequentialCell(
            nn.ReflectionPad2d((1, 1, 1, 1)),
            nn.Conv2d(features, features, 3, pad_mode='valid',
                      weight_init=init.Normal(0.2, 0.0)),
            norm_layer(features),
            nn.ReLU(),
            nn.ReflectionPad2d((1, 1, 1, 1)),
            nn.Conv2d(features, features, 3, pad_mode='valid',
                      weight_init=init.Normal(0.2, 0.0)),
            norm_layer(features)
        )

    def construct(self, x):
        return x + self.block(x)


##############################
#        Custom Layers
##############################


class AdaptiveInstanceNorm2d(nn.Cell):
    """Reference: https://github.com/NVlabs/MUNIT/blob/master/networks.py"""

    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__(AdaptiveInstanceNorm2d)
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        # weight and bias are dynamically assigned
        self.weight = init.initializer(init.Normal(0.2, 0.0), (256,))
        self.bias = init.initializer(init.Normal(0.02, 0.0), (256,))
        # just dummy buffers, not used
        self.running_var = Parameter(ops.ones(num_features), requires_grad=False)
        self.running_mean = Parameter(ops.zeros(num_features), requires_grad=False)

    def construct(self, x):
        assert (
                self.weight is not None and self.bias is not None
        ), "Please assign weight and bias before calling AdaIN!"
        b, c, h, w = x.shape
        running_mean = self.running_mean.tile((b,))
        running_var = self.running_var.tile((b,))

        # Apply instance norm
        x_reshaped = x.view(1, b * c, h, w)

        out = ops.batch_norm(
            x_reshaped, running_mean, running_var, self.weight, self.bias, True, self.momentum, self.eps
        )

        return out.view(b, c, h, w)

    def __repr__(self):
        return self.__class__.__name__ + "(" + str(self.num_features) + ")"


class LayerNorm(nn.Cell):
    """LayerNorm"""

    def __init__(self, num_features, eps=1e-5, affine=True):
        super().__init__(LayerNorm)
        self.num_features = num_features
        self.affine = affine
        self.eps = eps

        if self.affine:
            self.gamma = (Parameter(ops.rand(num_features)))
            self.beta = (Parameter(ops.zeros(num_features)))

    def construct(self, x):
        shape = [-1] + [1] * (x.dim() - 1)
        mean = x.view(x.shape[0], -1).mean(1).view(*shape)
        std = x.view(x.shape[0], -1).std(1).view(*shape)
        x = (x - mean) / (std + self.eps)

        if self.affine:
            shape = [1, -1] + [1] * (x.dim() - 2)
            x = x * self.gamma.view(*shape) + self.beta.view(*shape)
        return x

implementations/munit/models.py
mindspore-courses/MindSpore-GAN
"""MUNIT Model"""

import argparse
import datetime
import itertools
import os
import sys
import time

import mindspore
import mindspore.common.dtype as mstype
import numpy as np
from mindspore import nn, Tensor, ops
from mindspore.dataset.vision import transforms, Inter

from datasets import Edges2ShoesDataset
from img_utils import to_image
from models import Encoder, Decoder, MultiDiscriminator, LambdaLR

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="edges2shoes", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0001, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=128, help="size of image height")
parser.add_argument("--img_width", type=int, default=128, help="size of image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval saving generator samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between saving model checkpoints")
parser.add_argument("--n_downsample", type=int, default=2, help="number downsampling layers in encoder")
parser.add_argument("--n_residual", type=int, default=3, help="number of residual blocks in encoder / decoder")
parser.add_argument("--dim", type=int, default=64, help="number of filters in first encoder layer")
parser.add_argument("--style_dim", type=int, default=8, help="dimensionality of the style code")
opt = parser.parse_args()
print(opt)

# Create sample and checkpoint directories
os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

criterion_recon = nn.L1Loss()

# Initialize encoders, generators and discriminators
Enc1 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)
Dec1 = Decoder(dim=opt.dim, n_upsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)
Enc2 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)
Dec2 = Decoder(dim=opt.dim, n_upsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)
D1 = MultiDiscriminator()
D2 = MultiDiscriminator()

Enc1.update_parameters_name("Enc1")
Dec1.update_parameters_name("Dec1")
Enc2.update_parameters_name("Enc2")
Dec2.update_parameters_name("Dec2")
D1.update_parameters_name("D1")
D2.update_parameters_name("D2")

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/Enc1_{opt.epoch}.ckpt', Enc1)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/Dec1_{opt.epoch}.ckpt', Dec1)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/Enc2_{opt.epoch}.ckpt', Enc2)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/Dec2_{opt.epoch}.ckpt', Dec2)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D1_{opt.epoch}.ckpt', D1)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D2_{opt.epoch}.ckpt', D2)

# Loss weights
lambda_gan = 1
lambda_id = 10
lambda_style = 1
lambda_cont = 1
lambda_cyc = 0

# Learning rate update schedulers
decay_lr_G = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 49825)
decay_lr_D1 = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 49825)
decay_lr_D2 = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 49825)

# Optimizers
optimizer_G = nn.optim.Adam(itertools.chain(Enc1.trainable_params(), Dec1.trainable_params(),
                                            Enc2.trainable_params(), Dec2.trainable_params()),
                            learning_rate=decay_lr_G,
                            beta1=opt.b1, beta2=opt.b2)
optimizer_D1 = nn.optim.Adam(D1.trainable_params(), learning_rate=decay_lr_D1,
                             beta1=opt.b1, beta2=opt.b2)

optimizer_D2 = nn.optim.Adam(D2.trainable_params(), learning_rate=decay_lr_D2,
                             beta1=opt.b1, beta2=opt.b2)

# Dataset loader
transforms_ = [
    transforms.Resize((opt.img_height, opt.img_width), Inter.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
]

dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        transforms_=transforms_
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(opt.batch_size)

val_dataset = mindspore.dataset.GeneratorDataset(
    source=Edges2ShoesDataset(
        root="../../data/edges2shoes",
        transforms_=transforms_,
        mode='val'
    ),
    column_names=["A", "B"],
    shuffle=True
).batch(5)


def sample_image(batches):
    """Saves a generated sample from the validation set"""
    imgs = next(val_dataset.create_tuple_iterator())
    img_samples = None
    Enc1.set_train(False)
    Dec2.set_train(False)
    for img1, _ in zip(imgs[0], imgs[1]):
        # Create copies of image
        _X1 = img1.unsqueeze(0).tile((opt.style_dim, 1, 1, 1))
        # Get random style codes
        s_code = np.random.uniform(-1, 1, (opt.style_dim, opt.style_dim))
        s_code = Tensor(s_code, dtype=mstype.float32)
        # Generate samples
        c_code_1, _ = Enc1(_X1)
        _X12 = Dec2(c_code_1, s_code)
        # Concatenate samples horisontally
        _X12 = ops.cat(list(X12), -1)
        img_sample = ops.cat((img1, _X12), -1).unsqueeze(0)
        # Concatenate with previous samples vertically
        img_samples = img_sample if img_samples is None else ops.cat((img_samples, img_sample), -2)
    to_image(img_samples, os.path.join(f'images/{opt.dataset_name}', f'{batches}.png'))


def g_forward(_X1, _X2, _style_1, _style_2):
    """Generator warmup forward func"""
    # Get shared latent representation
    c_code_1, s_code_1 = Enc1(_X1)
    c_code_2, s_code_2 = Enc2(_X2)

    # Reconstruct images
    _X11 = Dec1(c_code_1, s_code_1)
    _X22 = Dec2(c_code_2, s_code_2)

    # Translate images
    _X21 = Dec1(c_code_2, _style_1)
    _X12 = Dec2(c_code_1, _style_2)

    # Cycle translation
    c_code_21, s_code_21 = Enc1(_X21)
    c_code_12, s_code_12 = Enc2(_X12)
    X121 = Dec1(c_code_12, s_code_1) if lambda_cyc > 0 else 0
    X212 = Dec2(c_code_21, s_code_2) if lambda_cyc > 0 else 0

    # Losses
    loss_GAN_1 = lambda_gan * D1.compute_loss(_X21, valid)
    loss_GAN_2 = lambda_gan * D2.compute_loss(_X12, valid)
    loss_ID_1 = lambda_id * criterion_recon(_X11, _X1)
    loss_ID_2 = lambda_id * criterion_recon(_X22, _X2)
    loss_s_1 = lambda_style * criterion_recon(s_code_21, _style_1)
    loss_s_2 = lambda_style * criterion_recon(s_code_12, _style_2)
    loss_c_1 = lambda_cont * criterion_recon(c_code_12, ops.stop_gradient(c_code_1))
    loss_c_2 = lambda_cont * criterion_recon(c_code_21, ops.stop_gradient(c_code_2))
    loss_cyc_1 = lambda_cyc * criterion_recon(X121, _X1) if lambda_cyc > 0 else 0
    loss_cyc_2 = lambda_cyc * criterion_recon(X212, _X2) if lambda_cyc > 0 else 0

    # Total loss
    _loss_G = (
            loss_GAN_1
            + loss_GAN_2
            + loss_ID_1
            + loss_ID_2
            + loss_s_1
            + loss_s_2
            + loss_c_1
            + loss_c_2
            + loss_cyc_1
            + loss_cyc_2
    )

    return _loss_G, _X12, _X21


def d1_forward(_X1, _X21, _valid, _fake):
    """Discriminator A forward function"""

    _loss_D1 = D1.compute_loss(_X1, _valid) + D1.compute_loss(_X21, _fake)

    return _loss_D1


def d2_forward(_X2, _X12, _valid, _fake):
    """Discriminator B forward function"""

    _loss_D2 = D2.compute_loss(_X2, _valid) + D2.compute_loss(_X12, _fake)

    return _loss_D2


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d1 = ops.value_and_grad(d1_forward, None, optimizer_D1.parameters, has_aux=False)
grad_d2 = ops.value_and_grad(d2_forward, None, optimizer_D2.parameters, has_aux=False)

# ----------
#  Training
# ----------

# Adversarial ground truths
valid = 1
fake = 0

prev_time = time.time()
for epoch in range(opt.n_epochs):
    # Model inputs
    for i, (X1, X2) in enumerate(dataset.create_tuple_iterator()):
        style_1 = ops.randn((X1.shape[0], opt.style_dim, 1, 1), dtype=mstype.float32)
        style_2 = ops.randn((X1.shape[0], opt.style_dim, 1, 1), dtype=mstype.float32)

        # -------------------------------
        #  Train Encoders and Generators
        # -------------------------------

        Enc1.set_train()
        Enc2.set_train()
        Dec1.set_train()
        Dec2.set_train()

        (loss_G, X12, X21), g_grads = grad_g(X1, X2, style_1, style_2)
        optimizer_G(g_grads)

        # -----------------------
        #  Train Discriminator 1
        # -----------------------

        D1.set_train()

        (loss_D1), d1_grads = grad_d1(X1, ops.stop_gradient(X21), valid, fake)
        optimizer_D1(d1_grads)

        # -----------------------
        #  Train Discriminator 2
        # -----------------------

        D2.set_train()

        (loss_D2), d2_grads = grad_d2(X2, ops.stop_gradient(X12), valid, fake)
        optimizer_D2(d2_grads)

        # --------------
        #  Log Progress
        # --------------

        # Determine approximate time left
        batches_done = epoch * dataset.get_dataset_size() + i
        batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {(loss_D1.asnumpy().item() + loss_D2.asnumpy().item()):.4f}] '
            f'[G loss: {loss_G.asnumpy().item():.4f}] '
            f'ETA: {time_left}'
        )

        if batches_done % opt.sample_interval == 0:
            sample_image(batches_done)

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        mindspore.save_checkpoint(Enc1, f'saved_models/{opt.dataset_name}/Enc1_{epoch}.ckpt')
        mindspore.save_checkpoint(Dec1, f'saved_models/{opt.dataset_name}/Dec1_{epoch}.ckpt')
        mindspore.save_checkpoint(Enc2, f'saved_models/{opt.dataset_name}/Enc2_{epoch}.ckpt')
        mindspore.save_checkpoint(Dec2, f'saved_models/{opt.dataset_name}/Dec2_{epoch}.ckpt')
        mindspore.save_checkpoint(D1, f'saved_models/{opt.dataset_name}/D1_{epoch}.ckpt')
        mindspore.save_checkpoint(D2, f'saved_models/{opt.dataset_name}/D2_{epoch}.ckpt')

implementations/munit/munit.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/pix2pix/img_utils.py
mindspore-courses/MindSpore-GAN
"""The Models of Pix2Pix"""

import mindspore.common.initializer as init
from mindspore import nn
from mindspore import ops


##############################
#           U-NET
##############################

class UNetDown(nn.Cell):
    """U-Net Down"""

    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
        super().__init__(UNetDown)
        layers = [
            nn.Conv2d(in_size, out_size, 4, 2,
                      'pad', 1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0))
        ]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_size, affine=False))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout(p=dropout))

        self.model = nn.SequentialCell(*layers)

    def construct(self, x):
        return self.model(x)


class UNetUp(nn.Cell):
    """U-Net Up"""

    def __init__(self, in_size, out_size, dropout=0.0):
        super().__init__(UNetUp)
        layers = [
            nn.Conv2dTranspose(in_size, out_size, 4, 2,
                               'pad', 1, has_bias=False,
                               weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(out_size, affine=False),
            nn.ReLU()
        ]
        if dropout:
            layers.append(nn.Dropout(p=dropout))

        self.model = nn.SequentialCell(*layers)

    def construct(self, x, skip_input):
        x = self.model(x)
        x = ops.cat((x, skip_input), 1)

        return x


class GeneratorUNet(nn.Cell):
    """U-Net Generator"""

    def __init__(self, in_channels=3, out_channels=3):
        super().__init__(GeneratorUNet)

        self.down1 = UNetDown(in_channels, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = UNetDown(512, 512, dropout=0.5)
        self.down7 = UNetDown(512, 512, dropout=0.5)
        self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5)

        self.up1 = UNetUp(512, 512, dropout=0.5)
        self.up2 = UNetUp(1024, 512, dropout=0.5)
        self.up3 = UNetUp(1024, 512, dropout=0.5)
        self.up4 = UNetUp(1024, 512, dropout=0.5)
        self.up5 = UNetUp(1024, 256)
        self.up6 = UNetUp(512, 128)
        self.up7 = UNetUp(256, 64)

        self.final = nn.SequentialCell(
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(128, out_channels, 4, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, x):
        # U-Net generator with skip connections from encoder to decoder
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        d7 = self.down7(d6)
        d8 = self.down8(d7)
        u1 = self.up1(d8, d7)
        u2 = self.up2(u1, d6)
        u3 = self.up3(u2, d5)
        u4 = self.up4(u3, d4)
        u5 = self.up5(u4, d3)
        u6 = self.up6(u5, d2)
        u7 = self.up7(u6, d1)

        return self.final(u7)


##############################
#        Discriminator
##############################

class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, in_channels=3):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, normalization=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4,
                                stride=2, pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.02, 0.0))]
            if normalization:
                layers.append(nn.InstanceNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *discriminator_block(in_channels * 2, 64, normalization=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            *discriminator_block(256, 512),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, 4,
                      pad_mode='pad', padding=1, has_bias=False)
        )

    def construct(self, img_A, img_B):
        # Concatenate image and condition image by channels to produce input
        img_input = ops.cat((img_A, img_B), 1)
        return self.model(img_input)

implementations/pix2pix/models.py
mindspore-courses/MindSpore-GAN
"""Pix2Pix Model"""

import argparse
import datetime
import os
import sys
import tarfile
import time
import urllib.request

import mindspore
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset import MindDataset
from mindspore.dataset.vision import transforms, Inter

from img_utils import to_image
from models import GeneratorUNet, Discriminator

file_path = "../../data/"

if not os.path.exists(file_path+"facades"):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path+"facades")
    url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/models/application/dataset_pix2pix.tar'
    print("Downloading facades dataset from " + url)
    urllib.request.urlretrieve(url, os.path.join(file_path, 'facades.tar'))
    tar = tarfile.open(os.path.join(file_path, 'facades.tar'), 'r')
    print("Unzipping dataset...")
    tar.extractall(path=file_path)
    tar.close()
    os.rename(file_path+'dataset_pix2pix',file_path+'facades')
    os.remove(os.path.join(file_path, 'facades.tar'))

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name",
                    type=str, default="facades", help="name of the dataset")
parser.add_argument("--batch_size",
                    type=int, default=1, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch",
                    type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height",
                    type=int, default=256, help="size of image height")
parser.add_argument("--img_width",
                    type=int, default=256, help="size of image width")
parser.add_argument("--channels",
                    type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval",
                    type=int, default=10, help="interval between sampling of images from generators"
                    )
parser.add_argument("--checkpoint_interval",
                    type=int, default=-1, help="interval between model checkpoints")
opt = parser.parse_args()
print(opt)

os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

# Loss functions
criterion_GAN = nn.MSELoss()
criterion_pixelwise = nn.L1Loss()

# Loss weight of L1 pixel-wise loss between translated image and real image
lambda_pixel = 100

# Calculate output of image discriminator (PatchGAN)
patch = (1, opt.img_height // 2 ** 4, opt.img_width // 2 ** 4)

# Initialize generator and discriminator
generator = GeneratorUNet()
discriminator = Discriminator()

optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

transform = [
    transforms.Resize((opt.img_height, opt.img_width), Inter.BICUBIC),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    transforms.HWC2CHW()
]

dataset = MindDataset("../../data/facades/train.mindrecord",
                      columns_list=["input_images", "target_images"], shuffle=True)

val_dataset = MindDataset("../../data/facades/train.mindrecord",
                          columns_list=["input_images", "target_images"], shuffle=True)


def sample_image(batches):
    """Saves a generated sample from the validation set"""
    imgs = next(val_dataset.create_dict_iterator())
    _input_images = Tensor(imgs["input_images"], dtype=mstype.float32)
    _target_images = Tensor(imgs["target_images"], dtype=mstype.float32)
    fake_images = generator(_input_images)
    img_sample = ops.cat((_input_images, fake_images, _target_images), -2)
    to_image(img_sample, os.path.join(f'images/{opt.dataset_name}', F'{batches}.png'))


def g_forward(_real_A, _real_B, _valid):
    """Generator forward function"""
    # GAN loss
    _fake_B = generator(_real_A)
    pred_fake = discriminator(_fake_B, _real_A)
    _loss_GAN = criterion_GAN(pred_fake, _valid)
    # Pixel-wise loss
    _loss_pixel = criterion_pixelwise(_fake_B, _real_B)
    # Total loss
    _g_loss = _loss_GAN + lambda_pixel * _loss_pixel
    return _g_loss, _fake_B, _loss_GAN, _loss_pixel


def d_forward(_real_A, _real_B, _fake_B, _valid, _fake):
    """Discriminator forward function"""
    # Real loss
    pred_real = discriminator(_real_B, _real_A)
    loss_real = criterion_GAN(pred_real, _valid)

    # Fake loss
    pred_fake = discriminator(_fake_B, _real_A)
    loss_fake = criterion_GAN(pred_fake, _fake)

    # Total loss
    _d_loss = 0.5 * (loss_real + loss_fake)
    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# ----------
#  Training
# ----------

prev_time = time.time()

generator.set_train()
discriminator.set_train()

for epoch in range(opt.n_epochs):
    # Model inputs
    for i, batch in enumerate(dataset.create_dict_iterator()):

        input_images = Tensor(batch["input_images"], dtype=mstype.float32)
        target_images = Tensor(batch["target_images"], dtype=mstype.float32)

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((input_images.shape[0], *patch)))
        fake = ops.stop_gradient(ops.zeros((input_images.shape[0], *patch)))

        (g_loss, fake_imgs, GAN_loss, pixel_loss), g_grads = grad_g(input_images, target_images, valid)
        optimizer_G(g_grads)

        (d_loss), d_grads = grad_d(input_images, target_images, ops.stop_gradient(fake_imgs),
                                   valid, fake)
        optimizer_D(d_grads)

        # --------------
        #  Log Progress
        # --------------
        # Determine approximate time left
        batches_done = epoch * dataset.get_dataset_size() + i
        batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        # Print log
        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G loss: {g_loss.asnumpy().item():.4f}, '
            f'pixel: {pixel_loss.asnumpy().item():.4f}, adv: {GAN_loss.asnumpy().item():.4f}] '
            f'ETA: {time_left}"'
        )

        # If at sample interval save image
        if batches_done % opt.sample_interval == 0:
            sample_image(batches_done)

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        mindspore.save_checkpoint(generator, f'saved_models/{opt.dataset_name}/generator_{epoch}.ckpt')
        mindspore.save_checkpoint(discriminator, f'saved_models/{opt.dataset_name}/discriminator_{epoch}.ckpt')

implementations/pix2pix/pix2pix.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/pixelda/img_utils.py
mindspore-courses/MindSpore-GAN
"""Dataset setting and data loader for MNIST-M.

Modified from
https://github.com/pytorch/vision/blob/master/torchvision/datasets/mnist.py

CREDIT: https://github.com/corenel
"""

from __future__ import print_function

import errno
import os
import gzip
import pickle

import mindspore
from PIL import Image
from mindspore import Tensor
import mindspore.common.dtype as mstype

# import essential packages
from six.moves import urllib


class MNISTM:
    """`MNIST-M Dataset."""

    url = "https://github.com/VanushVaswani/keras_mnistm/releases/download/1.0/keras_mnistm.pkl.gz"

    raw_folder = "raw"
    processed_folder = "processed"
    training_file = "mnist_m_train.pt"
    test_file = "mnist_m_test.pt"

    mnist_train_labels = []
    mnist_test_labels = []

    def __init__(self, root, mnist_root="data", train=True, transform=None, target_transform=None):
        """Init MNIST-M dataset."""
        super().__init__(MNISTM)
        self.root = os.path.expanduser(root)
        self.mnist_root = os.path.expanduser(mnist_root)
        self.transform = transform
        self.target_transform = target_transform
        self.train = train  # training set or test set

        self.train_ds, self.test_ds = self.dl_preprocess()

    def __getitem__(self, index):
        """Get images and target for data loader.

        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        if self.train:
            img, target = self.train_ds[0][index], self.train_ds[1][index]
        else:
            img, target = self.test_ds[0][index], self.test_ds[1][index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.squeeze(), mode="RGB")

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img[0], target

    def __len__(self):
        """Return size of dataset."""
        if self.train:
            return len(self.train_ds[0])
        return len(self.test_ds[0])

    def get_train_labels(self, labels):
        """Get MNIST labels"""
        for _, label in enumerate(labels.create_tuple_iterator()):
            self.mnist_train_labels.append(label)

    def get_test_labels(self, labels):
        """Get MNIST labels"""
        for _, label in enumerate(labels.create_tuple_iterator()):
            self.mnist_test_labels.append(label)

    def dl_preprocess(self):
        """Create MNIST-M dataset."""

        # make data dirs
        try:
            os.makedirs(os.path.join(self.root, self.raw_folder))
        except OSError as e:
            if e.errno == errno.EEXIST:
                pass
            else:
                raise

        # download pkl files
        filename = self.url.rpartition("/")[2]
        file_path = os.path.join(self.root, self.raw_folder, filename)
        if not os.path.exists(os.path.join(self.root, 'raw/keras_mnistm.pkl')):
            print("Downloading " + self.url)
            if not os.path.exists(file_path.replace(".gz", "")):
                data = urllib.request.urlopen(self.url)
                with open(file_path, "wb") as f:
                    f.write(data.read())
                with open(file_path.replace(".gz", ""), "wb") as out_f, gzip.GzipFile(file_path) as zip_f:
                    out_f.write(zip_f.read())
                os.unlink(file_path)

        # process and save as torch files
        print("Processing...")

        # load MNIST-M images from pkl file
        with open(file_path.replace(".gz", ""), "rb") as f:
            mnist_m_data = pickle.load(f, encoding="bytes")

        mnist_m_train_data = Tensor(mnist_m_data[b"train"], dtype=mstype.byte)
        mnist_m_test_data = Tensor(mnist_m_data[b"test"], dtype=mstype.byte)

        # get MNIST labels
        mnist_train = mindspore.dataset.MnistDataset(
            dataset_dir=self.mnist_root,
            usage='train',
            shuffle=False)
        mnist_test = mindspore.dataset.MnistDataset(
            dataset_dir=self.mnist_root,
            usage='test',
            shuffle=False)

        for _, (_, label) in enumerate(mnist_train.create_tuple_iterator()):
            self.mnist_train_labels.append(label)

        for _, (_, label) in enumerate(mnist_test.create_tuple_iterator()):
            self.mnist_test_labels.append(label)

        # save MNIST-M dataset
        training_set = (mnist_m_train_data.asnumpy(), self.mnist_train_labels)
        test_set = (mnist_m_test_data.asnumpy(), self.mnist_test_labels)

        print("Done!")
        return training_set, test_set

implementations/pixelda/mnistm.py
mindspore-courses/MindSpore-GAN
"""PixelDA Model"""

import argparse
import gzip
import itertools
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
import numpy as np
from mindspore import nn
from mindspore import ops
from mindspore.common import dtype as mstype
from mindspore.dataset.transforms import Compose
from mindspore.dataset.vision import transforms

from img_utils import to_image
from mnistm import MNISTM

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--n_residual_blocks", type=int, default=6, help="number of residual blocks in generator")
parser.add_argument("--latent_dim", type=int, default=10, help="dimensionality of the noise input")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--n_classes", type=int, default=10, help="number of classes in the dataset")
parser.add_argument("--sample_interval", type=int, default=300, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

# Calculate output of image discriminator (PatchGAN)
patch = int(opt.img_size / 2 ** 4)
patch = (1, patch, patch)


class ResidualBlock(nn.Cell):
    """Residual block"""

    def __init__(self, in_features=64):
        super().__init__(ResidualBlock)

        self.block = nn.SequentialCell(
            nn.Conv2d(in_features, in_features, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(in_features,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.ReLU(),
            nn.Conv2d(in_features, in_features, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(in_features,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False)
        )

    def construct(self, x):
        return x + self.block(x)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        # Fully-connected layer which constructs image channel shaped output from noise
        self.fc = nn.Dense(opt.latent_dim, opt.channels * opt.img_size ** 2)

        self.l1 = nn.SequentialCell(
            nn.Conv2d(opt.channels * 2, 64, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.ReLU()
        )

        resblocks = []
        for _ in range(opt.n_residual_blocks):
            resblocks.append(ResidualBlock())
        self.resblocks = nn.SequentialCell(*resblocks)

        self.l2 = nn.SequentialCell(
            nn.Conv2d(64, opt.channels, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, img, z):
        gen_input = ops.cat((img, self.fc(z).view(*img.shape)), 1)
        out = self.l1(gen_input)
        out = self.resblocks(out)
        img_ = self.l2(out)

        return img_


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def block(in_features, out_features, normalization=True):
            """Discriminator block"""
            layers = [nn.Conv2d(in_features, out_features, 3,
                                stride=2, pad_mode='pad', padding=1,
                                weight_init=init.Normal(0.02, 0.0)),
                      nn.LeakyReLU(0.2)
                      ]
            if normalization:
                layers.append(nn.BatchNorm2d(out_features))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.channels, 64, normalization=False),
            *block(64, 128),
            *block(128, 256),
            *block(256, 512),
            nn.Conv2d(512, 1, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0))
        )

    def construct(self, img):
        validity = self.model(img)

        return validity


class Classifier(nn.Cell):
    """Classifier Network"""
    def __init__(self):
        super().__init__(Classifier)

        def block(in_features, out_features, normalization=True):
            """Classifier block"""
            layers = [
                nn.Conv2d(in_features, out_features, 3, stride=2,
                          pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2)
            ]
            if normalization:
                layers.append(nn.BatchNorm2d(out_features,
                                             gamma_init=init.Normal(0.02, 1.0),
                                             beta_init=init.Constant(0.0), affine=False))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.channels, 64, normalization=False), *block(64, 128), *block(128, 256), *block(256, 512)
        )

        input_size = opt.img_size // 2 ** 4
        self.output_layer = nn.SequentialCell(
            nn.Dense(512 * input_size ** 2, opt.n_classes),
            nn.Softmax()
        )

    def construct(self, img):
        feature_repr = self.model(img)
        feature_repr = feature_repr.view(feature_repr.shape[0], -1)
        label = self.output_layer(feature_repr)
        return label


# Loss function
adversarial_loss = nn.MSELoss()
task_loss = nn.CrossEntropyLoss()

# Loss weights
lambda_adv = 1
lambda_task = 0.1

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()
classifier = Classifier()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

os.makedirs("../../data/MNIST-M", exist_ok=True)

dataset1 = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

dataset2 = mindspore.dataset.GeneratorDataset(
    source=MNISTM(
        root='../../data/MNIST-M',
        mnist_root='../../data/MNIST',
        transform=Compose(
            [
                transforms.Resize(opt.img_size),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
            ]
        )
    ),
    shuffle=True,
    column_names=["image", "target"]
).batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(itertools.chain(generator.trainable_params(), classifier.trainable_params()),
                            learning_rate=opt.lr,
                            beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1,
                            beta2=opt.b2)


def g_forward(_imgs_A, _labels_A, _valid):
    """Generator forward function"""
    # Sample noise as generator input
    z = ops.randn((batch_size, opt.latent_dim))

    # Generate a batch of images
    _fake_B = generator(_imgs_A, z)

    # Perform task on translated source image
    _label_pred = classifier(_fake_B)

    # Calculate the task loss
    task_loss_ = (task_loss(_label_pred, _labels_A) + task_loss(classifier(_imgs_A), _labels_A)) / 2

    # Loss measures generator's ability to fool the discriminator
    _g_loss = lambda_adv * adversarial_loss(discriminator(_fake_B), _valid) + lambda_task * task_loss_

    return _g_loss, _fake_B, _label_pred


def d_forward(_imgs_B, _fake_B, _valid, _fake):
    """Discriminator forward function"""
    # Measure discriminator's ability to classify real from generated samples
    real_loss = adversarial_loss(discriminator(_imgs_B), _valid)
    fake_loss = adversarial_loss(discriminator(_fake_B), _fake)
    _d_loss = (real_loss + fake_loss) / 2

    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

generator.set_train()
discriminator.set_train()
classifier.set_train()

# ----------
#  Training
# ----------

# Keeps 100 accuracy measurements
task_performance = []
target_performance = []

for epoch in range(opt.n_epochs):
    for i, ((imgs_A, labels_A), (imgs_B, labels_B)) in enumerate(zip(dataset1.create_tuple_iterator(),
                                                                     dataset2.create_tuple_iterator())):
        batch_size = imgs_A.shape[0]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((batch_size, *patch)))
        fake = ops.stop_gradient(ops.zeros((batch_size, *patch)))

        labels_A = ops.Cast()(labels_A, mstype.int32)
        labels_B = ops.Cast()(labels_B, mstype.int32)

        # Configure input
        imgs_A = imgs_A.broadcast_to((batch_size, 3, opt.img_size, opt.img_size))

        # ------------------
        #  Train Generators
        # ------------------

        (g_loss, fake_B, label_pred), g_grads = grad_g(imgs_A, labels_A, valid)
        optimizer_G(g_grads)

        # ----------------------
        #  Train Discriminators
        # ----------------------

        (d_loss), d_grads = grad_d(imgs_B, ops.stop_gradient(fake_B), valid, fake)
        optimizer_D(d_grads)

        # ---------------------------------------
        #  Evaluate Performance on target domain
        # ---------------------------------------

        # Evaluate performance on translated Domain A
        acc = np.mean(np.argmax(label_pred.asnumpy(), axis=1) == labels_A.asnumpy())
        task_performance.append(acc)
        if len(task_performance) > 100:
            task_performance.pop(0)

        # Evaluate performance on Domain B
        pred_B = classifier(imgs_B)
        target_acc = np.mean(np.argmax(pred_B.asnumpy(), axis=1) == labels_B.asnumpy())
        target_performance.append(target_acc)
        if len(target_performance) > 100:
            target_performance.pop(0)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset1.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G loss: {g_loss.asnumpy().item():.4f}] '
            f'[CLF acc: {100 * acc:3f}% ({100 * np.mean(task_performance):3}%), '
            f'target_acc: {100 * target_acc:3}% ({100 * np.mean(target_performance):3}%)]'
        )

        batches_done = epoch * dataset1.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            sample = ops.cat((imgs_A[:5], fake_B[:5], imgs_B[:5]), -2)
            to_image(sample, os.path.join("images", F'{batches_done}.png'))

implementations/pixelda/pixelda.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/relativistic_gan/img_utils.py
mindspore-courses/MindSpore-GAN
"""RelativisticGAN Model"""

import argparse
import os

import mindspore
from mindspore import nn
from mindspore import ops
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling")
parser.add_argument("--rel_avg_gan", action="store_true", help="relativistic average GAN instead of standard")
opt = parser.parse_args()
print(opt)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.init_size = opt.img_size // 4
        self.l1 = nn.SequentialCell(nn.Dense(opt.latent_dim, 128 * self.init_size ** 2))

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3, stride=1,
                      pad_mode='pad', padding=1),
            nn.BatchNorm2d(128, 0.8, affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3, stride=1,
                      pad_mode='pad', padding=1),
            nn.BatchNorm2d(64, 0.8, affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3, stride=1,
                      pad_mode='pad', padding=1),
            nn.Tanh()
        )

    def construct(self, _z):
        out = self.l1(_z)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [
                nn.Conv2d(in_filters, out_filters, 3, 2,
                          'pad', 1),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8, affine=False))
            return block

        self.model = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4
        self.adv_layer = nn.SequentialCell(nn.Dense(128 * ds_size ** 2, 1))

    def construct(self, _img):
        out = self.model(_img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)

        return validity


# Loss function
adversarial_loss = nn.BCEWithLogitsLoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_real_imgs, _valid):
    """Generator forward function"""
    # Generate a batch of images
    z = ops.randn((_real_imgs.shape[0], opt.latent_dim), dtype=mstype.float32)
    _gen_imgs = generator(z)

    real_pred = ops.stop_gradient(discriminator(_real_imgs))
    fake_pred = discriminator(_gen_imgs)

    if opt.rel_avg_gan:
        _g_loss = adversarial_loss(fake_pred - real_pred.mean(0, keep_dims=True), _valid)
    else:
        _g_loss = adversarial_loss(fake_pred - real_pred, _valid)

    # Loss measures generator's ability to fool the discriminator
    _g_loss = adversarial_loss(discriminator(_gen_imgs), _valid)

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _valid, _fake):
    """Discriminator forward function"""

    # Predict validity
    real_pred = discriminator(_real_imgs)
    fake_pred = discriminator(_gen_imgs)

    if opt.rel_avg_gan:
        real_loss = adversarial_loss(real_pred - fake_pred.mean(0, keep_dims=True), _valid)
        fake_loss = adversarial_loss(fake_pred - real_pred.mean(0, keep_dims=True), _fake)
    else:
        real_loss = adversarial_loss(real_pred - fake_pred, _valid)
        fake_loss = adversarial_loss(fake_pred - real_pred, _fake)

    _d_loss = (real_loss + fake_loss) / 2

    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# ----------
#  Training
# ----------

generator.set_train()
discriminator.set_train()

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs.shape[0], 1)))
        fake = ops.stop_gradient(ops.zeros((imgs.shape[0], 1)))

        # Configure input
        real_imgs = imgs

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(real_imgs, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs), valid, fake)
        optimizer_D(d_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/relativistic_gan/relativistic_gan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/sgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""SGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import mindspore.common.initializer as init
import numpy as np
from mindspore import nn, Tensor
from mindspore import ops
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--num_classes", type=int, default=10, help="number of classes for dataset")
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        self.label_emb = nn.Embedding(opt.num_classes, opt.latent_dim)

        self.init_size = opt.img_size // 4  # Initial size before upsampling
        self.l1 = nn.SequentialCell(
            nn.Dense(opt.latent_dim, 128 * self.init_size ** 2)
        )

        self.conv_blocks = nn.SequentialCell(
            nn.BatchNorm2d(128,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 128, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(128, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2.0, recompute_scale_factor=True),
            nn.Conv2d(128, 64, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, 0.8,
                           gamma_init=init.Normal(0.02, 1.0),
                           beta_init=init.Constant(0.0), affine=False),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, opt.channels, 3,
                      stride=1, pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        )

    def construct(self, noise):
        out = self.l1(noise)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        def discriminator_block(in_filters, out_filters, bn=True):
            """Returns layers of each discriminator block"""
            block = [
                nn.Conv2d(in_filters, out_filters, 3,
                          stride=2, pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8,
                                            gamma_init=init.Normal(0.02, 1.0),
                                            beta_init=init.Constant(0.0), affine=False))
            return block

        self.conv_blocks = nn.SequentialCell(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = opt.img_size // 2 ** 4

        # Output layers
        self.adv_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, 1),
            nn.Sigmoid()
        )
        self.aux_layer = nn.SequentialCell(
            nn.Dense(128 * ds_size ** 2, opt.num_classes + 1),
            nn.Softmax()
        )

    def construct(self, img):
        out = self.conv_blocks(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        label = self.aux_layer(out)

        return validity, label


# Loss functions
adversarial_loss = nn.BCELoss()
auxiliary_loss = nn.CrossEntropyLoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def g_forward(_batch_size, _valid):
    """Generator forward function"""
    # Sample noise and labels as generator input
    z = ops.randn((_batch_size, opt.latent_dim), dtype=mstype.float32)

    # Generate a batch of images
    _gen_imgs = generator(z)

    # Loss measures generator's ability to fool the discriminator
    validity, _ = discriminator(_gen_imgs)
    _g_loss = adversarial_loss(validity, _valid)

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _labels, _fake_aux_gt, _valid, _fake):
    """Discriminator forward function"""
    # Loss for real images
    real_pred, real_aux = discriminator(_real_imgs)
    d_real_loss = (adversarial_loss(real_pred, _valid) + auxiliary_loss(real_aux, _labels)) / 2

    # Loss for fake images
    fake_pred, fake_aux = discriminator(_gen_imgs)
    d_fake_loss = (adversarial_loss(fake_pred, _fake) + auxiliary_loss(fake_aux, _fake_aux_gt)) / 2

    # Total discriminator loss
    _d_loss = (d_real_loss + d_fake_loss) / 2

    # Calculate discriminator accuracy
    pred = np.concatenate([real_aux.asnumpy(), fake_aux.asnumpy()], axis=0)
    gt = np.concatenate([_labels.asnumpy(), _fake_aux_gt.asnumpy()], axis=0)
    _d_acc = np.mean(np.argmax(pred, axis=1) == gt)

    return _d_loss, _d_acc


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=True)

generator.set_train()
discriminator.set_train()

# ----------
#  Training
# ----------

for epoch in range(opt.n_epochs):
    for i, (imgs, labels) in enumerate(dataset.create_tuple_iterator()):

        batch_size = imgs.shape[0]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((batch_size, 1)))
        fake = ops.stop_gradient(ops.zeros((batch_size, 1)))
        fake_aux_gt = ops.stop_gradient(ops.fill(mstype.int32, Tensor(batch_size), opt.num_classes))

        # Configure input
        real_imgs = imgs
        labels = Tensor(labels, dtype=mstype.int32)

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(batch_size, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------
        (d_loss, d_acc), d_grads = grad_d(real_imgs, ops.stop_gradient(gen_imgs),
                                          labels, fake_aux_gt, valid, fake)
        optimizer_D(d_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}, '
            f'acc: {100 * d_acc}%] '
            f'[G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/sgan/sgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/softmax_gan/img_utils.py
mindspore-courses/MindSpore-GAN
"""SoftmaxGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from mindspore import nn
from mindspore import ops
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("Downloading MNIST dataset from" + url)
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("Unzipping...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, _z):
        img = self.model(_z)
        img = img.view(img.shape[0], *img_shape)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        self.model = nn.SequentialCell(
            nn.Dense(opt.img_size ** 2, 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dense(256, 1),
        )

    def construct(self, _img):
        img_flat = _img.view(_img.shape[0], -1)
        validity = self.model(img_flat)

        return validity


# Loss functions
adversarial_loss = nn.BCELoss()

# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()

transform = [
    transforms.Resize(opt.img_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)


def log(x):
    """Log function"""
    return ops.log(x + 1e-8)


def g_forward(_real_imgs, _g_target):
    """Generator forward function"""
    # Generate a batch of images
    z = ops.randn((_real_imgs.shape[0], opt.latent_dim), dtype=mstype.float32)
    _gen_imgs = generator(z)

    _d_real = discriminator(_real_imgs)
    _d_fake = discriminator(_gen_imgs)

    # Partition function
    Z = ops.sum(ops.exp(ops.neg(_d_real))) + ops.sum(ops.exp(ops.neg(_d_fake)))

    # Calculate loss of generator and update
    _g_loss = _g_target * (ops.sum(_d_real) + ops.sum(_d_fake)) + log(Z)

    return _g_loss, _gen_imgs


def d_forward(_real_imgs, _gen_imgs, _d_target):
    """Discriminator forward function"""

    _d_real = discriminator(_real_imgs)
    _d_fake = discriminator(_gen_imgs)

    # Partition function
    Z = ops.sum(ops.exp(ops.neg(_d_real))) + ops.sum(ops.exp(ops.neg(_d_fake)))

    # Calculate loss of discriminator and update
    _d_loss = _d_target * ops.sum(_d_real) + log(Z)

    return _d_loss, _d_real, _d_fake


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=True)

# ----------
#  Training
# ----------

generator.set_train()
discriminator.set_train()

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):

        batch_size = imgs.shape[0]

        # Adversarial ground truths
        g_target = 1 / (batch_size * 2)
        d_target = 1 / batch_size

        # Configure input
        real_imgs = imgs

        # -----------------
        #  Train Generator
        # -----------------

        (g_loss, gen_imgs), g_grads = grad_g(real_imgs, g_target)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss, d_real, d_fake), d_grads = grad_d(real_imgs,
                                                   ops.stop_gradient(gen_imgs), d_target)
        optimizer_D(d_grads)

        # --------------
        # Log Progress
        # --------------

        print(
            f'[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G loss: {g_loss.asnumpy().item():.4f}]'
        )
        batches_done = epoch * dataset.get_dataset_size() + i
        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))

implementations/softmax_gan/softmax_gan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/srgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""The Models of SRGAN"""

from mindcv.models import vgg19
from mindspore import nn
from mindspore import ops


class FeatureExtractor(nn.Cell):
    """VGG Feature Extractor"""

    def __init__(self):
        super().__init__(FeatureExtractor)
        vgg19_model = vgg19(pretrained=True)
        self.feature_extractor = nn.SequentialCell(
            *list(vgg19_model.features.cell_list[:18])
        )

    def construct(self, img):
        return self.feature_extractor(img)


class ResidualBlock(nn.Cell):
    """Residual Block"""

    def __init__(self, in_features):
        super().__init__(ResidualBlock)
        self.conv_block = nn.SequentialCell(
            nn.Conv2d(in_features, in_features,
                      3, 1, 'pad', 1),
            nn.BatchNorm2d(in_features, 0.8),
            nn.PReLU(),
            nn.Conv2d(in_features, in_features,
                      3, 1, 'pad', 1),
            nn.BatchNorm2d(in_features, 0.8),
        )

    def construct(self, x):
        return x + self.conv_block(x)


class GeneratorResNet(nn.Cell):
    """ResNet Generator"""

    def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16):
        super().__init__(GeneratorResNet)

        # First layer
        self.conv1 = nn.SequentialCell(
            nn.Conv2d(in_channels, 64,
                      9, 1, 'pad', 4),
            nn.PReLU()
        )

        # Residual blocks
        res_blocks = []
        for _ in range(n_residual_blocks):
            res_blocks.append(ResidualBlock(64))
        self.res_blocks = nn.SequentialCell(*res_blocks)

        # Second conv layer post residual blocks
        self.conv2 = nn.SequentialCell(
            nn.Conv2d(64, 64,
                      3, 1, 'pad', 1),
            nn.BatchNorm2d(64, 0.8)
        )

        # Upsampling layers
        upsampling = []
        for _ in range(2):
            upsampling += [
                # nn.Upsample(scale_factor=2),
                nn.Conv2d(64, 256,
                          3, 1, 'pad', 1),
                nn.BatchNorm2d(256, affine=False),
                nn.PixelShuffle(upscale_factor=2),
                nn.PReLU()
            ]
        self.upsampling = nn.SequentialCell(*upsampling)

        # Final output layer
        self.conv3 = nn.SequentialCell(
            nn.Conv2d(64, out_channels,
                      9, 1, 'pad', 4),
            nn.Tanh()
        )

    def construct(self, x):
        out1 = self.conv1(x)
        out = self.res_blocks(out1)
        out2 = self.conv2(out)
        out = ops.add(out1, out2)
        out = self.upsampling(out)
        out = self.conv3(out)
        return out


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self, input_shape):
        super().__init__(Discriminator)

        self.input_shape = input_shape
        in_channels, in_height, in_width = self.input_shape
        patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)
        self.output_shape = (1, patch_h, patch_w)

        def discriminator_block(in_filters, out_filters, first_block=False):
            layers = []
            layers.append(nn.Conv2d(in_filters, out_filters, 3,
                                    1, 'pad', 1))
            if not first_block:
                layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            layers.append(nn.Conv2d(out_filters, out_filters, 3,
                                    2, 'pad', 1))
            layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        layers = []
        in_filters = in_channels
        for i, out_filters in enumerate([64, 128, 256, 512]):
            layers.extend(discriminator_block(in_filters, out_filters, first_block=i == 0))
            in_filters = out_filters

        layers.append(nn.Conv2d(out_filters, 1, 3,
                                1, 'pad', 1))

        self.model = nn.SequentialCell(*layers)

    def construct(self, img):
        return self.model(img)

implementations/srgan/models.py
mindspore-courses/MindSpore-GAN
"""
Super-resolution of CelebA using Generative Adversarial Networks.
The dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0
(if not available there see if options are listed at http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)
Instrustion on running the script:
1. Download the dataset from the provided link
2. Save the folder 'img_align_celeba' to '../../data/'
4. Run the sript using command 'python3 srgan.py'
"""

import argparse
import os
import sys

import mindspore
from mindspore import nn
from mindspore import ops
from mindspore.dataset import CelebADataset
from mindspore.dataset.vision import transforms, Inter

from img_utils import to_image, make_grid
from models import GeneratorResNet, Discriminator, FeatureExtractor

os.makedirs("images", exist_ok=True)
os.makedirs("saved_models", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=4, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--hr_height", type=int, default=256, help="high res. image height")
parser.add_argument("--hr_width", type=int, default=256, help="high res. image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=100, help="interval between saving image samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
opt = parser.parse_args()
print(opt)

hr_shape = (opt.hr_height, opt.hr_width)

# Initialize generator and discriminator
generator = GeneratorResNet()
discriminator = Discriminator(input_shape=(opt.channels, *hr_shape))
feature_extractor = FeatureExtractor()

# Set feature extractor to inference mode
feature_extractor.set_train(False)

criterion_GAN = nn.MSELoss()
criterion_content = nn.L1Loss()

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/generator_{opt.epoch}.ckpt', generator)
    mindspore.load_checkpoint(f'saved_models/discriminator{opt.epoch}.ckpt', discriminator)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), opt.lr, opt.b1, opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), opt.lr, opt.b1, opt.b2)


def preprocess(_imgs):
    """Dataset preprocess func"""
    _imgs = transforms.Resize((opt.hr_height, opt.hr_width), Inter.BICUBIC)(_imgs)
    _imgs_lr = transforms.Resize((opt.hr_height // 4, opt.hr_width // 4), Inter.BICUBIC)(_imgs)
    _imgs = transforms.ToTensor()(_imgs)
    _imgs = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], is_hwc=False)(_imgs)
    _imgs_lr = transforms.ToTensor()(_imgs_lr)
    _imgs_lr = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], is_hwc=False)(_imgs_lr)

    return _imgs, _imgs_lr


dataset = CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True
).map(operations=[preprocess], input_columns=["image"], output_columns=["image_hr", "image_lr"]).batch(opt.batch_size)


def g_forward(_imgs_hr, _imgs_lr, _valid):
    """Generator forward function"""
    # Generate a high resolution image from low resolution input
    _gen_hr = generator(_imgs_lr)

    # Adversarial loss
    loss_GAN = criterion_GAN(discriminator(_gen_hr), _valid)

    # Content loss
    gen_features = feature_extractor(_gen_hr)
    real_features = feature_extractor(_imgs_hr)
    loss_content = criterion_content(gen_features, ops.stop_gradient(real_features))

    # Total loss
    _g_loss = loss_content + 1e-3 * loss_GAN
    return _g_loss, _gen_hr


def d_forward(_imgs_hr, _gen_hr, _valid, _fake):
    """Discriminator forward function"""
    # Loss of real and fake images
    loss_real = criterion_GAN(discriminator(_imgs_hr), _valid)
    loss_fake = criterion_GAN(discriminator(_gen_hr), _fake)

    # Total loss
    _d_loss = (loss_real + loss_fake) / 2
    return _d_loss


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=False)

# ----------
#  Training
# ----------

saved_samples = {}

for epoch in range(opt.n_epochs):
    # Model inputs
    for i, batch in enumerate(dataset.create_dict_iterator()):
        imgs_hr = batch["image_hr"]
        imgs_lr = batch["image_lr"]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((imgs_lr.shape[0], *discriminator.output_shape)))
        fake = ops.stop_gradient(ops.zeros((imgs_lr.shape[0], *discriminator.output_shape)))

        # ------------------
        #  Train Generators
        # ------------------
        (g_loss, gen_hr), g_grads = grad_g(imgs_hr, imgs_lr, valid)
        optimizer_G(g_grads)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss), d_grads = grad_d(imgs_hr, ops.stop_gradient(gen_hr),
                                   valid, fake)
        optimizer_D(d_grads)

        # --------------
        #  Log Progress
        # --------------

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {d_loss.asnumpy().item():.4f}] '
            f'[G loss: {g_loss.asnumpy().item():.4f}]'
        )

        batches_done = epoch * dataset.get_dataset_size() + i
        # If at sample interval save image
        if batches_done % opt.sample_interval == 0:
            # Save image grid with upsampled inputs and SRGAN outputs
            imgs_lr = ops.interpolate(imgs_lr, scale_factor=4.0, recompute_scale_factor=True)
            gen_hr = make_grid(gen_hr, nrow=1, normalize=True)
            imgs_lr = make_grid(imgs_lr, nrow=1, normalize=True)
            img_grid = ops.cat((imgs_lr, gen_hr), -1)
            to_image(img_grid, f'images/{batches_done}.png')

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        mindspore.save_checkpoint(generator, f'saved_models/{opt.dataset_name}/generator_{epoch}.ckpt')
        mindspore.save_checkpoint(discriminator, f'saved_models/{opt.dataset_name}/discriminator_{epoch}.ckpt')

implementations/srgan/srgan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/stargan/img_utils.py
mindspore-courses/MindSpore-GAN
"""The Models of StarGAN"""

import mindspore.common.initializer as init
from mindspore import nn, ops


##############################
#           RESNET
##############################

class ResidualBlock(nn.Cell):
    """Residual block"""
    def __init__(self, in_features):
        super().__init__(ResidualBlock)

        conv_block = [
            nn.Conv2d(in_features, in_features, 3, 1, 'pad', 1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(in_features, affine=True),
            nn.ReLU(),
            nn.Conv2d(in_features, in_features, 3, 1, 'pad', 1, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(in_features, affine=True),
        ]

        self.conv_block = nn.SequentialCell(*conv_block)

    def construct(self, x):
        return x + self.conv_block(x)


class GeneratorResNet(nn.Cell):
    """Generator Network"""
    def __init__(self, img_shape=(3, 128, 128), res_blocks=9, c_dim=40):
        super().__init__(GeneratorResNet)
        channels, _, _ = img_shape

        # Initial convolution block
        model = [
            nn.Conv2d(channels + c_dim, 64, 7, 1, 'pad', 3, has_bias=False,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.InstanceNorm2d(64, affine=True),
            nn.ReLU()
        ]

        # Downsampling
        curr_dim = 64
        for _ in range(2):
            model += [
                nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 'pad', 1, has_bias=False,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.InstanceNorm2d(curr_dim * 2, affine=True),
                nn.ReLU()
            ]
            curr_dim *= 2

        # Residual blocks
        for _ in range(res_blocks):
            model += [ResidualBlock(curr_dim)]

        # Upsampling
        for _ in range(2):
            model += [
                nn.Conv2dTranspose(curr_dim, curr_dim // 2, 4, 2, 'pad', 1, has_bias=False,
                                   weight_init=init.Normal(0.02, 0.0)),
                nn.InstanceNorm2d(curr_dim // 2, affine=True),
                nn.ReLU()
            ]
            curr_dim = curr_dim // 2

        # Output layer
        model += [
            nn.Conv2d(curr_dim, channels, 7, 1, 'pad', 3,
                      weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        ]

        self.model = nn.SequentialCell(*model)

    def construct(self, x, c):
        c = c.view(c.shape[0], c.shape[1], 1, 1)
        c = c.tile((1, 1, x.shape[2], x.shape[3]))
        x = ops.cat((x, c), 1)
        return self.model(x)


##############################
#        Discriminator
##############################

class Discriminator(nn.Cell):
    """Discriminator Network"""
    def __init__(self, img_shape=(3, 128, 128), c_dim=40, n_strided=6):
        super().__init__(Discriminator)
        channels, img_size, _ = img_shape

        def discriminator_block(in_filters, out_filters):
            """Returns downsampling layers of each discriminator block"""
            layers = [
                nn.Conv2d(in_filters, out_filters, 4, 2, 'pad', 1),
                nn.LeakyReLU(0.01)]
            return layers

        layers = discriminator_block(channels, 64)
        curr_dim = 64
        for _ in range(n_strided - 1):
            layers.extend(discriminator_block(curr_dim, curr_dim * 2))
            curr_dim *= 2

        self.layers = nn.SequentialCell(*layers)

        # Output 1: PatchGAN
        self.out1 = nn.Conv2d(curr_dim, 1, 3, pad_mode='pad', padding=1, has_bias=False,
                              weight_init=init.Normal(0.02, 0.0))
        # Output 2: Class prediction
        kernel_size = img_size // 2 ** n_strided
        self.out2 = nn.Conv2d(curr_dim, c_dim, kernel_size, has_bias=False,
                              weight_init=init.Normal(0.02, 0.0),
                              pad_mode='valid')

    def construct(self, img):
        feature_repr = self.layers(img)
        out_adv = self.out1(feature_repr)
        out_cls = self.out2(feature_repr)
        return out_adv, out_cls.view(out_cls.shape[0], -1)

implementations/stargan/models.py
mindspore-courses/MindSpore-GAN
"""
StarGAN (CelebA)
The dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0
And the annotations: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AAA8YmAHNNU6BEfWMPMfM6r9a/Anno?dl=0&preview=list_attr_celeba.txt
Instructions on running the script:
1. Download the dataset and annotations from the provided link
2. Copy 'list_attr_celeba.txt' to folder 'img_align_celeba'
2. Save the folder 'img_align_celeba' to '../../data/'
4. Run the script by 'python3 stargan.py'
"""

import argparse
import datetime
import os
import sys
import time

import mindspore
import mindspore.common.dtype as mstype
from mindspore import nn
from mindspore import ops
from mindspore.dataset import CelebADataset
from mindspore.dataset.vision import transforms, Inter

from img_utils import to_image
from models import GeneratorResNet, Discriminator

os.makedirs("images", exist_ok=True)
os.makedirs("saved_models", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=16, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=128, help="size of image height")
parser.add_argument("--img_width", type=int, default=128, help="size of image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval between saving generator samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
parser.add_argument("--residual_blocks", type=int, default=6, help="number of residual blocks in generator")
parser.add_argument(
    "--selected_attrs",
    "--list",
    nargs="+",
    help="selected attributes for the CelebA dataset",
    default=["Black_Hair", "Blond_Hair", "Brown_Hair", "Male", "Young"],
)
parser.add_argument("--n_critic", type=int, default=5, help="number of training iterations for WGAN discriminator")
opt = parser.parse_args()
print(opt)

mindspore.set_context(mode=0)

c_dim = 40
img_shape = (opt.channels, opt.img_height, opt.img_width)

criterion_cycle = nn.L1Loss()


def criterion_cls(logit, target):
    """cls criterion"""
    return (ops.binary_cross_entropy_with_logits(logit, target, weight=ops.ones_like(logit), reduction='mean',
                                                 pos_weight=ops.ones_like(logit)) / logit.shape[0])


# Loss weights
lambda_cls = 1
lambda_rec = 10
lambda_gp = 10

# Initialize generator and discriminator
generator = GeneratorResNet(img_shape=img_shape, res_blocks=opt.residual_blocks, c_dim=c_dim)

discriminator = Discriminator(img_shape=img_shape, c_dim=c_dim)

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/generator_{opt.epoch}.ckpt', generator)
    mindspore.load_checkpoint(f'saved_models/discriminator{opt.epoch}.ckpt', discriminator)

# Optimizers
optimizer_G = nn.optim.Adam(generator.trainable_params(), opt.lr, opt.b1, opt.b2)
optimizer_D = nn.optim.Adam(discriminator.trainable_params(), opt.lr, opt.b1, opt.b2)

train_transforms = [
    transforms.Resize(int(1.12 * opt.img_height), Inter.BICUBIC),
    transforms.RandomCrop(opt.img_height),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
]

val_transforms = [
    transforms.Resize((opt.img_height, opt.img_width), Inter.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), is_hwc=False),
]

train_dataset = (CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True,
).map(operations=train_transforms, input_columns=["image"])
                 .batch(opt.batch_size))

val_dataset = CelebADataset(
    dataset_dir="../../data/CelebA",
    shuffle=True,
    decode=True
).map(operations=val_transforms, input_columns=["image"]).batch(10)


def compute_gradient_penalty(real_samples, fake_samples):
    """Calculates the gradient penalty loss for WGAN GP"""
    # Random weight term for interpolation between real and fake samples
    alpha = ops.rand((real_samples.shape[0], 1, 1, 1))
    # Get random interpolation between real and fake samples
    interpolates = alpha * real_samples + ((1 - alpha) * fake_samples)
    # Get gradient w.r.t. interpolates
    grad_fn = ops.grad(discriminator)
    gradients = grad_fn(interpolates)
    gradients = gradients.view(gradients.shape[0], -1)
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty


label_changes = [
    ((0, 1), (1, 0), (2, 0)),  # Set to black hair
    ((0, 0), (1, 1), (2, 0)),  # Set to blonde hair
    ((0, 0), (1, 0), (2, 1)),  # Set to brown hair
    ((3, -1),),  # Flip gender
    ((4, -1),),  # Age flip
]


def sample_images(_batches_done):
    """Saves a generated sample of domain translations"""
    val_imgs, val_labels = next(val_dataset.create_tuple_iterator())

    img_samples = None
    for j in range(10):
        img, label = val_imgs[j], val_labels[j]
        # Repeat for number of label changes
        _imgs = img.tile((c_dim, 1, 1, 1))
        _labels = label.tile((c_dim, 1))
        # Make changes to labels
        for sample_i, changes in enumerate(label_changes):
            for col, val in changes:
                if val == -1:
                    _labels[sample_i, col] = 1 - _labels[sample_i, col].asnumpy().item()
                else:
                    _labels[sample_i, col] = val

        # Generate translations
        _labels = ops.Cast()(_labels, mstype.float32)
        gen_imgs = generator(_imgs, _labels)
        # Concatenate images by width
        gen_imgs = ops.cat(list(gen_imgs), -1)
        img_sample = ops.cat((img, gen_imgs), -1)
        # Add as row to generated samples
        img_samples = img_sample if img_samples is None else ops.cat((img_samples, img_sample), -2)

    to_image(img_samples.view(1, *img_samples.shape), f'images/{_batches_done}.png')


def g_forward(_imgs, _labels, _sampled_c):
    """Generator warmup forward func"""
    # Translate and reconstruct image
    gen_imgs = generator(_imgs, _sampled_c)
    recov_imgs = generator(gen_imgs, _labels)
    # Discriminator evaluates translated image
    fake_validity, pred_cls = discriminator(gen_imgs)
    # Adversarial loss
    _g_loss_adv = -ops.mean(fake_validity)
    # Classification loss
    _g_loss_cls = criterion_cls(pred_cls, _sampled_c)
    # Reconstruction loss
    _g_loss_rec = criterion_cycle(recov_imgs, _imgs)
    # Total loss
    _g_loss = _g_loss_adv + lambda_cls * _g_loss_cls + lambda_rec * _g_loss_rec

    return _g_loss, _g_loss_adv, _g_loss_cls, _g_loss_rec


def d_forward(_imgs, _fake_imgs, _labels):
    """Discriminator forward function"""
    # Real images
    real_validity, pred_cls = discriminator(_imgs)
    pred_cls = ops.Cast()(pred_cls, mstype.float32)
    # Fake images
    fake_validity, _ = discriminator(_fake_imgs)
    # Gradient penalty
    gradient_penalty = compute_gradient_penalty(_imgs, _fake_imgs)
    # Adversarial loss
    _d_loss_adv = -ops.mean(real_validity) + ops.mean(fake_validity) + lambda_gp * gradient_penalty
    # Classification loss
    _d_loss_cls = criterion_cls(pred_cls, _labels)
    # Total loss
    _d_loss = _d_loss_adv + lambda_cls * _d_loss_cls
    return _d_loss, _d_loss_adv, _d_loss_cls


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, optimizer_D.parameters, has_aux=True)

# ----------
#  Training
# ----------

generator.set_train()
discriminator.set_train()

saved_samples = []
start_time = time.time()

for epoch in range(opt.epoch, opt.n_epochs):
    # Model inputs
    for i, (imgs, labels) in enumerate(train_dataset.create_tuple_iterator()):
        # Sample labels as generator inputs
        sampled_c = ops.randint(0, 2, (imgs.shape[0], c_dim))
        sampled_c = ops.Cast()(sampled_c, mstype.float32)
        labels = ops.Cast()(labels, mstype.float32)
        # Generate fake batch of images
        fake_imgs = generator(imgs, sampled_c)

        # ---------------------
        #  Train Discriminator
        # ---------------------

        (d_loss, d_loss_adv, d_loss_cls), d_grads = grad_d(imgs, ops.stop_gradient(fake_imgs), labels)
        optimizer_D(d_grads)

        # Every n_critic times update generator
        if i % opt.n_critic == 0:
            # -----------------
            #  Train Generator
            # -----------------

            (g_loss, g_loss_adv, g_loss_cls, g_loss_rec), g_grads = grad_g(imgs, labels, sampled_c)
            optimizer_G(g_grads)

            # --------------
            #  Log Progress
            # --------------

            # Determine approximate time left
            batches_done = epoch * train_dataset.get_dataset_size() + i
            batches_left = opt.n_epochs * train_dataset.get_dataset_size() - batches_done
            time_left = datetime.timedelta(seconds=batches_left * (time.time() - start_time) / (batches_done + 1))

            # Print log
            sys.stdout.write(
                f'\r[Epoch {epoch}/{opt.n_epochs}] '
                f'[Batch {i}/{train_dataset.get_dataset_size()}] '
                f'[D adv: {d_loss_adv.asnumpy().item():.4f}, '
                f'aux: {d_loss_cls.asnumpy().item():.4f}] '
                f'[G loss: {g_loss.asnumpy().item():.4f}, '
                f'adv: {g_loss_adv.asnumpy().item():.4f}, '
                f'aux: {g_loss_cls.asnumpy().item():.4f}, '
                f'cycle: {g_loss_rec.asnumpy().item():.4f}, '
                f'ETA: {time_left}"'
            )

            # If at sample interval sample and save image
            if batches_done % opt.sample_interval == 0:
                sample_images(batches_done)

        if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
            # Save model checkpoints
            mindspore.save_checkpoint(generator, f'saved_models/generator_{epoch}.ckpt')
            mindspore.save_checkpoint(discriminator, f'saved_models/discriminator_{epoch}.ckpt')

implementations/stargan/stargan.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image


def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/unit/img_utils.py
mindspore-courses/MindSpore-GAN
"""UNIT Networks"""

import mindspore.common.dtype as mstype
import mindspore.common.initializer as init
from mindspore import nn, ops, Tensor
from mindspore.nn.learning_rate_schedule import LearningRateSchedule


class LambdaLR(LearningRateSchedule):
    """Learning Rate Schedule"""

    def __init__(self, lr, n_epochs, offset, decay_start_epoch, step_per_epoch):
        super().__init__()
        assert (n_epochs - decay_start_epoch) > 0, "Decay must start before the training session ends!"
        self.lr = lr
        self.n_epochs = n_epochs
        self.offset = offset
        self.decay_start_epoch = decay_start_epoch
        self.cast = ops.Cast()
        self.step_per_epoch = step_per_epoch

    def construct(self, global_step):
        epoch = self.cast(global_step, mstype.float32) // self.step_per_epoch
        return self.lr * (1.0 - max(Tensor(0.0), epoch + self.offset - self.decay_start_epoch) / (
                self.n_epochs - self.decay_start_epoch))


##############################
#           RESNET
##############################

class ResidualBlock(nn.Cell):
    """Residual block"""

    def __init__(self, features):
        super().__init__(ResidualBlock)

        conv_block = [
            nn.ReflectionPad2d(1),
            nn.Conv2d(features, features, 3, pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(features, affine=False),
            nn.ReLU(),
            nn.ReflectionPad2d(1),
            nn.Conv2d(features, features, 3, pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(features, affine=False),
        ]

        self.conv_block = nn.SequentialCell(*conv_block)

    def construct(self, x):
        return x + self.conv_block(x)


class Encoder(nn.Cell):
    """"Encoder"""

    def __init__(self, in_channels=3, dim=64, n_downsample=2, shared_block=None):
        super().__init__(Encoder)

        # Initial convolution block
        layers = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(in_channels, dim, 7, pad_mode='pad',
                      weight_init=init.Normal(0.02, 0.0)),
            nn.BatchNorm2d(64, affine=False),
            nn.LeakyReLU(0.2),
        ]

        # Downsampling
        for _ in range(n_downsample):
            layers += [
                nn.Conv2d(dim, dim * 2, 4, stride=2,
                          pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0)),
                nn.BatchNorm2d(dim * 2, affine=False),
                nn.ReLU()
            ]
            dim *= 2

        # Residual blocks
        for _ in range(3):
            layers += [ResidualBlock(dim)]

        self.model_blocks = nn.SequentialCell(*layers)
        self.shared_block = shared_block

    def reparameterization(self, _mu):
        """Reparameterization func"""
        _z = ops.randn(_mu.shape, dtype=mstype.float32)
        return _z + _mu

    def construct(self, x):
        x = self.model_blocks(x)
        mu = self.shared_block(x)
        z = self.reparameterization(mu)
        return mu, z


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self, out_channels=3, dim=64, n_upsample=2, shared_block=None):
        super().__init__(Generator)

        self.shared_block = shared_block

        layers = []
        dim = dim * 2 ** n_upsample
        # Residual blocks
        for _ in range(3):
            layers += [ResidualBlock(dim)]

        # Upsampling
        for _ in range(n_upsample):
            layers += [
                nn.Conv2dTranspose(dim, dim // 2, 4, stride=2,
                                   pad_mode='pad', padding=1,
                                   weight_init=init.Normal(0.02, 0.0)),
                nn.BatchNorm2d(dim // 2, affine=False),
                nn.LeakyReLU(0.2)
            ]
            dim = dim // 2

        # Output layer
        layers += [
            nn.ReflectionPad2d(3),
            nn.Conv2d(dim, out_channels, 7,
                      pad_mode="pad", weight_init=init.Normal(0.02, 0.0)),
            nn.Tanh()
        ]

        self.model_blocks = nn.SequentialCell(*layers)

    def construct(self, x):
        x = self.shared_block(x)
        x = self.model_blocks(x)
        return x


##############################
#        Discriminator
##############################


class Discriminator(nn.Cell):
    """Discriminator"""

    def __init__(self, input_shape):
        super().__init__(Discriminator)
        channels, height, width = input_shape
        # Calculate output of image discriminator (PatchGAN)
        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)

        def discriminator_block(in_filters, out_filters, normalize=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [
                nn.Conv2d(in_filters, out_filters, 4, stride=2,
                          pad_mode='pad', padding=1,
                          weight_init=init.Normal(0.02, 0.0))
            ]
            if normalize:
                layers.append(nn.BatchNorm2d(out_filters, affine=False))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *discriminator_block(channels, 64, normalize=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            *discriminator_block(256, 512),
            nn.Conv2d(512, 1, 3,
                      pad_mode='pad', padding=1,
                      weight_init=init.Normal(0.02, 0.0))
        )

    def construct(self, img):
        return self.model(img)

implementations/unit/models.py
mindspore-courses/MindSpore-GAN
"""UNIT Model"""

import argparse
import datetime
import itertools
import os
import sys
import time
import urllib
import zipfile

import mindspore
from mindspore import nn, ops
from mindspore.dataset import MindDataset
from mindspore.nn import CellList

from img_utils import to_image
from models import Encoder, ResidualBlock, Generator, LambdaLR, Discriminator

file_path = "../../data/"

if not os.path.exists(file_path + "apple2orange"):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path + "apple2orange")
    url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/models/application/CycleGAN_apple2orange.zip'
    print("Downloading apple2orange dataset from " + url)
    urllib.request.urlretrieve(url, os.path.join(file_path, 'apple2orange.zip'))

    while not os.path.exists(os.path.join(file_path, 'apple2orange.zip')):
        pass

    zippedData = zipfile.ZipFile(os.path.join(file_path, 'apple2orange.zip'), 'r')
    print("Unzipping dataset...")
    zippedData.extractall(path=file_path)
    zippedData.close()
    os.rename(file_path + 'CycleGAN_apple2orange', file_path + 'apple2orange')
    os.remove(os.path.join(file_path, 'apple2orange.zip'))

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=200, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="apple2orange", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0001, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=256, help="size of image height")
parser.add_argument("--img_width", type=int, default=256, help="size of image width")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=1, help="interval between saving generator samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between saving model checkpoints")
parser.add_argument("--n_downsample", type=int, default=2, help="number downsampling layers in encoder")
parser.add_argument("--dim", type=int, default=64, help="number of filters in first encoder layer")
opt = parser.parse_args()
print(opt)

# Create sample and checkpoint directories
os.makedirs(f'images/{opt.dataset_name}', exist_ok=True)
os.makedirs(f'saved_models/{opt.dataset_name}', exist_ok=True)

# Losses
criterion_GAN = nn.MSELoss()
criterion_pixel = nn.L1Loss()

input_shape = (opt.channels, opt.img_height, opt.img_width)

# Dimensionality (channel-wise) of image embedding
shared_dim = opt.dim * 2 ** opt.n_downsample

# Initialize generator and discriminator
shared_E = ResidualBlock(features=shared_dim)
E1 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, shared_block=shared_E)
E2 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, shared_block=shared_E)
shared_G = ResidualBlock(features=shared_dim)
G1 = Generator(dim=opt.dim, n_upsample=opt.n_downsample, shared_block=shared_G)
G2 = Generator(dim=opt.dim, n_upsample=opt.n_downsample, shared_block=shared_G)
D1 = Discriminator(input_shape)
D2 = Discriminator(input_shape)

shared_E.update_parameters_name("shared_E_")
shared_G.update_parameters_name("shared_G_")
E1.update_parameters_name("E1_")
E2.update_parameters_name("E2_")
G1.update_parameters_name("G1_")
G2.update_parameters_name("G2_")
D1.update_parameters_name("D1_")
D2.update_parameters_name("D2_")

if opt.epoch != 0:
    # Load pretrained models
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/E1_{opt.epoch}.ckpt', E1)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/E2_{opt.epoch}.ckpt', E2)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G1_{opt.epoch}.ckpt', G1)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/G2_{opt.epoch}.ckpt', G2)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D1_{opt.epoch}.ckpt', D1)
    mindspore.load_checkpoint(f'saved_models/%{opt.dataset_name}/D2_{opt.epoch}.ckpt', D2)

# Loss weights
lambda_0 = 10  # GAN
lambda_1 = 0.1  # KL (encoded images)
lambda_2 = 100  # ID pixel-wise
lambda_3 = 0.1  # KL (encoded translated images)
lambda_4 = 100  # Cycle pixel-wise

# Learning rate update schedulers
decay_lr_G_AB = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 500)
decay_lr_G_BA = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 500)
decay_lr_D_A = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 500)
decay_lr_D_B = LambdaLR(opt.lr, opt.n_epochs, opt.epoch, opt.decay_epoch, 500)

no_shared_cell_E1 = CellList()
no_shared_cell_G1 = CellList()
for name, cell in E1.cells_and_names():
    if "shared_block" not in name and name != '':
        no_shared_cell_E1.append(cell)

for name, cell in G1.cells_and_names():
    if "shared_block" not in name and name != '':
        no_shared_cell_G1.append(cell)

no_shared_cell_E1.update_parameters_name("E1")
no_shared_cell_G1.update_parameters_name("G1")

optimizer_G = nn.optim.Adam(itertools.chain(no_shared_cell_E1.trainable_params(),
                                            E2.trainable_params(),
                                            no_shared_cell_G1.trainable_params(),
                                            G2.trainable_params()),
                            learning_rate=decay_lr_G_AB,
                            beta1=opt.b1, beta2=opt.b2)

optimizer_D1 = nn.optim.Adam(D1.trainable_params(), learning_rate=decay_lr_D_A,
                             beta1=opt.b1, beta2=opt.b2)

optimizer_D2 = nn.optim.Adam(D2.trainable_params(), learning_rate=decay_lr_D_B,
                             beta1=opt.b1, beta2=opt.b2)

dataset = (MindDataset("../../data/apple2orange/apple2orange_train.mindrecord", shuffle=True).
           batch(opt.batch_size))

val_dataset = (MindDataset("../../data/apple2orange/apple2orange_train.mindrecord", shuffle=True).
               batch(5))


def sample_images(batches):
    """Saves a generated sample from the test set"""
    imgs = next(val_dataset.create_dict_iterator())
    G1.set_train(False)
    G2.set_train(False)
    _X1 = imgs["image_A"]
    _X2 = imgs["image_B"]

    _, Z1 = E1(_X1)
    _, Z2 = E2(_X2)
    _fake_X1 = G1(Z2)
    _fake_X2 = G2(Z1)

    img_sample = ops.cat((_X1, _fake_X2, _X2, _fake_X1), 0)
    to_image(img_sample, os.path.join(f'images/{opt.dataset_name}', F'{batches}.png'))


def compute_kl(mu):
    """Compute KL"""
    mu_2 = ops.pow(mu, 2)
    loss = ops.mean(mu_2)
    return loss


def g_forward(_X1, _X2, _valid):
    """Generator forward function"""
    # Get shared latent representation
    mu1, Z1 = E1(_X1)
    mu2, Z2 = E2(_X2)

    # Reconstruct images
    recon_X1 = G1(Z1)
    recon_X2 = G2(Z2)

    # Translate images
    _fake_X1 = G1(Z2)
    _fake_X2 = G2(Z1)

    # Cycle translation
    mu1_, Z1_ = E1(_fake_X1)
    mu2_, Z2_ = E2(_fake_X2)
    cycle_X1 = G1(Z2_)
    cycle_X2 = G2(Z1_)

    # Losses
    loss_GAN_1 = lambda_0 * criterion_GAN(D1(_fake_X1), _valid)
    loss_GAN_2 = lambda_0 * criterion_GAN(D2(_fake_X2), _valid)
    loss_KL_1 = lambda_1 * compute_kl(mu1)
    loss_KL_2 = lambda_1 * compute_kl(mu2)
    loss_ID_1 = lambda_2 * criterion_pixel(recon_X1, _X1)
    loss_ID_2 = lambda_2 * criterion_pixel(recon_X2, _X2)
    loss_KL_1_ = lambda_3 * compute_kl(mu1_)
    loss_KL_2_ = lambda_3 * compute_kl(mu2_)
    loss_cyc_1 = lambda_4 * criterion_pixel(cycle_X1, _X1)
    loss_cyc_2 = lambda_4 * criterion_pixel(cycle_X2, _X2)

    # Total loss
    _loss_G = (
            loss_KL_1
            + loss_KL_2
            + loss_ID_1
            + loss_ID_2
            + loss_GAN_1
            + loss_GAN_2
            + loss_KL_1_
            + loss_KL_2_
            + loss_cyc_1
            + loss_cyc_2
    )

    return _loss_G, _fake_X1, _fake_X2


def d1_forward(_X1, _fake_X1, _valid, _fake):
    """Discriminator forward function"""

    _loss_D1 = criterion_GAN(D1(_X1), _valid) + criterion_GAN(D1(_fake_X1), _fake)

    return _loss_D1


def d2_forward(_X2, _fake_X2, _valid, _fake):
    """Discriminator forward function"""

    _loss_D2 = criterion_GAN(D2(X2), _valid) + criterion_GAN(D2(_fake_X2), _fake)

    return _loss_D2


grad_g = ops.value_and_grad(g_forward, None, optimizer_G.parameters, has_aux=True)
grad_d1 = ops.value_and_grad(d1_forward, None, optimizer_D1.parameters)
grad_d2 = ops.value_and_grad(d2_forward, None, optimizer_D2.parameters)

# ----------
#  Training
# ----------

G1.set_train()
G2.set_train()
D1.set_train()
D2.set_train()
E1.set_train()
E2.set_train()

prev_time = time.time()
for epoch in range(opt.n_epochs):
    # Model inputs
    for i, batch in enumerate(dataset.create_dict_iterator()):
        # Set model input
        X1 = batch["image_A"]
        X2 = batch["image_B"]

        # Adversarial ground truths
        valid = ops.stop_gradient(ops.ones((X1.shape[0], *D1.output_shape)))
        fake = ops.stop_gradient(ops.zeros((X1.shape[0], *D1.output_shape)))

        # -------------------------------
        #  Train Encoders and Generators
        # -------------------------------

        (loss_G, fake_X1, fake_X2), g_grads = grad_g(X1, X2, valid)
        optimizer_G(g_grads)

        # -----------------------
        #  Train Discriminator 1
        # -----------------------

        (loss_D1), d1_grads = grad_d1(X1, ops.stop_gradient(fake_X1), valid, fake)
        optimizer_D1(d1_grads)

        # -----------------------
        #  Train Discriminator 2
        # -----------------------

        (loss_D2), d2_grads = grad_d2(X2, ops.stop_gradient(fake_X2), valid, fake)
        optimizer_D2(d2_grads)

        # --------------
        #  Log Progress
        # --------------

        # Determine approximate time left
        batches_done = epoch * dataset.get_dataset_size() + i
        batches_left = opt.n_epochs * dataset.get_dataset_size() - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        sys.stdout.write(
            f'\r[Epoch {epoch}/{opt.n_epochs}] '
            f'[Batch {i}/{dataset.get_dataset_size()}] '
            f'[D loss: {(loss_D1.asnumpy().item() + loss_D2.asnumpy().item()):.4f}] '
            f'[G loss: {loss_G.asnumpy().item():.4f}] '
            f'ETA: {time_left}'
        )

        # If at sample interval save image
        if batches_done % opt.sample_interval == 0:
            sample_images(batches_done)

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        mindspore.save_checkpoint(E1, f'saved_models/{opt.dataset_name}/E1_{epoch}.ckpt')
        mindspore.save_checkpoint(E2, f'saved_models/{opt.dataset_name}/E2_{epoch}.ckpt')
        mindspore.save_checkpoint(G1, f'saved_models/{opt.dataset_name}/G1_{epoch}.ckpt')
        mindspore.save_checkpoint(G2, f'saved_models/{opt.dataset_name}/G2_{epoch}.ckpt')
        mindspore.save_checkpoint(D1, f'saved_models/{opt.dataset_name}/D1_{epoch}.ckpt')
        mindspore.save_checkpoint(D2, f'saved_models/{opt.dataset_name}/D2_{epoch}.ckpt')

implementations/unit/unit.py
mindspore-courses/MindSpore-GAN
"""图像处理"""
import math
import pathlib
from typing import BinaryIO, Union, Optional, Tuple
import mindspore
from mindspore import ops
import mindspore.numpy as mnp
from PIL import Image



def make_grid(
        tensor,
        nrow: int = 5,
        padding: int = 2,
        normalize: bool = True,
        value_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: float = 0.0):
    """grid"""
    if isinstance(tensor, list):
        tensor = mnp.stack(tensor, axis=0)

    if tensor.ndim == 2:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 3:
        if tensor.shape[0] == 1:
            tensor = mnp.concatenate((tensor, tensor, tensor), 0)
        tensor = tensor.expand_dims(0)

    if tensor.ndim == 4 and tensor.shape[1] == 1:
        tensor = mnp.concatenate((tensor, tensor, tensor), 1)

    if normalize is True:
        tensor_list = []
        if value_range is not None:
            assert isinstance(
                value_range, tuple
            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"

        def norm_ip(img, low, high):
            img = ops.clip_by_value(img, low, high)
            img = (img - low) / (max(high - low, 1e-5))
            return img

        def norm_range(input_tensor, value_range):
            if value_range is not None:
                return norm_ip(input_tensor, value_range[0], value_range[1])
            return norm_ip(input_tensor, float(input_tensor.min()), float(input_tensor.max()))

        if scale_each is True:
            for input_t in tensor:  # loop over mini-batch dimension
                tensor_list.append(norm_range(input_t, value_range))
        else:
            tensor_list = norm_range(tensor, value_range)

        if isinstance(tensor_list, mindspore.Tensor):
            tensor = tensor_list
        else:
            tensor = mnp.concatenate(tensor_list)

    assert isinstance(tensor, mindspore.Tensor)
    if tensor.shape[0] == 1:
        return tensor.squeeze(0)

    # make the mini-batch of images into a grid
    nmaps = tensor.shape[0]
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.shape[2] + padding), int(tensor.shape[3] + padding)
    num_channels = tensor.shape[1]
    grid = mnp.full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            grid[:, y * height + padding: (y + 1) * height, x * width + padding: (x + 1) * width] = tensor[k]
            k = k + 1
    return grid


def to_image(tensor,
             file: Union[str, pathlib.Path, BinaryIO],
             _format=None,
             **kwargs):
    """保存图片"""
    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid * 255 + 0.5
    ndarr = ops.clip_by_value(ndarr, 0, 255).transpose(1, 2, 0).astype(mindspore.uint8).asnumpy()
    img = Image.fromarray(ndarr)
    img.save(file, format=_format)

implementations/wgan/img_utils.py
mindspore-courses/MindSpore-GAN
"""WGAN Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

from img_utils import to_image

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.00005, help="learning rate")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size",
                    type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--n_critic",
                    type=int, default=5, help="number of training steps for discriminator per iter")
parser.add_argument("--clip_value",
                    type=float, default=0.01, help="lower and upper clip value for disc. weights")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, _z):
        img = self.model(_z)
        img = img.view(img.shape[0], *img_shape)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        self.model = nn.SequentialCell(
            nn.Dense(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dense(256, 1)
        )

    def construct(self, img):
        img_flat = img.view(img.shape[0], -1)
        validity = self.model(img_flat)
        return validity


generator = Generator()
discriminator = Discriminator()

G_Optim = nn.optim.RMSProp(generator.trainable_params(), learning_rate=opt.lr, decay=0.99, epsilon=1e-08)
D_Optim = nn.optim.RMSProp(discriminator.trainable_params(), learning_rate=opt.lr, decay=0.99, epsilon=1e-08)


def g_forward(_z):
    """Generator forward function"""
    _gen_imgs = generator(_z)
    _loss_G = -ops.mean(discriminator(_gen_imgs))
    return _loss_G, _gen_imgs


# 判别器正向传播
def d_forward(_real_imgs, _z):
    """Discriminator forward function"""
    fake_imgs = ops.stop_gradient(generator(_z))
    real = discriminator(_real_imgs)
    fake = discriminator(fake_imgs)
    _loss_D = -ops.mean(real) + ops.mean(fake)

    return _loss_D


transform = [
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5], is_hwc=False)
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, G_Optim.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, D_Optim.parameters, has_aux=False)

batches_done = 0
for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        # Train Discriminator
        discriminator.set_train()
        generator.set_train(False)
        real_imgs = Tensor(imgs, dtype=mstype.float32)

        z = ops.randn((imgs.shape[0], opt.latent_dim))

        (loss_D), D_grads = grad_d(real_imgs, z)
        D_Optim(D_grads)

        # Clip weights of discriminator
        for p in discriminator.get_parameters():
            ops.assign(p, ops.clamp(p, -opt.clip_value, opt.clip_value))

        # Train the generator every n_critic iterations
        if i % opt.n_critic == 0:
            generator.set_train()
            (loss_G, gen_imgs), G_grads = grad_g(z)
            G_Optim(G_grads)

            print(
                f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
                f'[D loss: {loss_D.asnumpy().item():.4f}] [G loss: {loss_G.asnumpy().item():.4f}]'
            )

        if batches_done % opt.sample_interval == 0:
            to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))
        batches_done += 1

implementations/wgan/wgan.py
mindspore-courses/MindSpore-GAN
"""WGAN DIV model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from matplotlib import pyplot as plt
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=64, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size",
                    type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--n_critic",
                    type=int, default=5, help="number of training steps for discriminator per iter")
parser.add_argument("--clip_value",
                    type=float, default=0.01, help="lower and upper clip value for disc. weights")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)

image_path = "./images"


def save_imgs(_gen_imgs, idx):
    """Save the generated test image."""
    for j in range(_gen_imgs.shape[0]):
        plt.subplot(5, 5, j + 1)
        plt.imshow(_gen_imgs[j, 0, :, :] / 2 + 0.5, cmap="gray")
        plt.axis("off")
    plt.savefig(f'image_path/test_{idx}.png"')


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, _z):
        img = self.model(_z)
        img = img.view(img.shape[0], *img_shape)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        self.model = nn.SequentialCell(
            nn.Dense(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dense(256, 1)
        )

    def construct(self, img):
        img_flat = img.view(img.shape[0], -1)
        validity = self.model(img_flat)
        return validity


k = 2
p = 6

generator = Generator()
discriminator = Discriminator()
generator.update_parameters_name('generator')
discriminator.update_parameters_name('discriminator')
generator.set_train()
discriminator.set_train()

G_Optim = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
D_Optim = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

G_Optim.update_parameters_name('G_Optim')
D_Optim.update_parameters_name('D_Optim')


def g_forward(_z):
    """Generator forward function"""
    fake_imgs = generator(_z)
    # Loss measures generator's ability to fool the discriminator
    # Train on fake images
    fake_validity = discriminator(fake_imgs)
    g_loss = -ops.mean(fake_validity)
    return g_loss, fake_imgs


def d_forward(_real_imgs):
    """Discriminator forward function"""
    _z = ops.randn((imgs.shape[0], opt.latent_dim))
    fake_imgs = generator(_z)
    # Real images
    real_validity = discriminator(_real_imgs)
    # Fake images
    fake_validity = discriminator(fake_imgs)

    # Compute W-div gradient penalty
    real_grad_out = ops.ones((_real_imgs.shape[0], 1))
    ops.stop_gradient(real_grad_out)
    real_grad_fn = ops.grad(discriminator)
    real_gradients = real_grad_fn(_real_imgs)
    real_grad_norm = real_gradients.view(real_gradients.shape[0], -1).pow(2).sum(1) ** (p / 2)

    fake_grad_out = ops.ones((_real_imgs.shape[0], 1))
    ops.stop_gradient(fake_grad_out)
    fake_grad_fn = ops.grad(discriminator)
    fake_gradients = fake_grad_fn(fake_imgs)
    fake_grad_norm = fake_gradients.view(fake_gradients.shape[0], -1).pow(2).sum(1) ** (p / 2)

    div_gp = ops.mean(real_grad_norm + fake_grad_norm) * k / 2
    # Adversarial loss
    d_loss = -ops.mean(real_validity) + ops.mean(fake_validity) + div_gp

    return d_loss, _z


transform = [
    transforms.Rescale(1.0 / 255.0, 0),
    transforms.Resize(opt.img_size),
    transforms.Normalize([0.5], [0.5]),
    transforms.HWC2CHW()
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, G_Optim.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, D_Optim.parameters, has_aux=True)

batches_done = 0

generator.update_parameters_name('generator')
discriminator.update_parameters_name('discriminator')
discriminator.set_train()
generator.set_train()

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        real_imgs = Tensor(imgs, dtype=mstype.float32)

        (loss_D, z), D_grads = grad_d(real_imgs)
        D_Optim(D_grads)

        # Train the generator every n_critic iterations
        if i % opt.n_critic == 0:
            (loss_G, gen_imgs), G_grads = grad_g(z)
            G_Optim(G_grads)

            print(
                f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
                f'[D loss: {loss_D.asnumpy().item()}] [G loss: {loss_G.asnumpy().item()}]'
            )

            if batches_done % opt.sample_interval == 0:
                save_imgs(gen_imgs[:25].asnumpy(), batches_done)
                # to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))
            batches_done += opt.n_critic

implementations/wgan_div/wgan_div.py
mindspore-courses/MindSpore-GAN
"""WGAN GP Model"""

import argparse
import gzip
import os
import shutil
import urllib.request

import mindspore
import numpy as np
from matplotlib import pyplot as plt
from mindspore import Tensor, ops
from mindspore import nn
from mindspore.common import dtype as mstype
from mindspore.dataset.vision import transforms

file_path = "../../data/MNIST/"

if not os.path.exists(file_path):
    # 下载数据集
    if not os.path.exists('../../data'):
        os.mkdir('../../data')
    os.mkdir(file_path)
    base_url = 'http://yann.lecun.com/exdb/mnist/'
    file_names = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
                  't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
    for file_name in file_names:
        url = (base_url + file_name).format(**locals())
        print("正在从" + url + "下载MNIST数据集...")
        urllib.request.urlretrieve(url, os.path.join(file_path, file_name))
        with gzip.open(os.path.join(file_path, file_name), 'rb') as f_in:
            print("正在解压数据集...")
            with open(os.path.join(file_path, file_name)[:-3], 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(os.path.join(file_path, file_name))

os.makedirs("images", exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs",
                    type=int, default=200, help="number of epochs of training")
parser.add_argument("--batch_size",
                    type=int, default=128, help="size of the batches")
parser.add_argument("--lr",
                    type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1",
                    type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
                    type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_cpu",
                    type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim",
                    type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size",
                    type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels",
                    type=int, default=1, help="number of image channels")
parser.add_argument("--n_critic",
                    type=int, default=5, help="number of training steps for discriminator per iter")
parser.add_argument("--clip_value",
                    type=float, default=0.01, help="lower and upper clip value for disc. weights")
parser.add_argument("--sample_interval",
                    type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)

img_shape = (opt.channels, opt.img_size, opt.img_size)

image_path = "./images"


def save_imgs(_gen_imgs, idx):
    """Save the generated test image."""
    for j in range(_gen_imgs.shape[0]):
        plt.subplot(5, 5, j + 1)
        plt.imshow(_gen_imgs[j, 0, :, :] / 2 + 0.5, cmap="gray")
        plt.axis("off")
    plt.savefig(f'image_path/test_{idx}.png"')


class Generator(nn.Cell):
    """Generator Network"""

    def __init__(self):
        super().__init__(Generator)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Dense(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.SequentialCell(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Dense(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def construct(self, _z):
        img = self.model(_z)
        img = img.view(img.shape[0], *img_shape)
        return img


class Discriminator(nn.Cell):
    """Discriminator Network"""

    def __init__(self):
        super().__init__(Discriminator)

        self.model = nn.SequentialCell(
            nn.Dense(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2),
            nn.Dense(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dense(256, 1)
        )

    def construct(self, img):
        img_flat = img.view(img.shape[0], -1)
        validity = self.model(img_flat)
        return validity


generator = Generator()
discriminator = Discriminator()
generator.update_parameters_name('generator')
discriminator.update_parameters_name('discriminator')
generator.set_train()
discriminator.set_train()

# Loss weight for gradient penalty
lambda_gp = 10

G_Optim = nn.optim.Adam(generator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)
D_Optim = nn.optim.Adam(discriminator.trainable_params(), learning_rate=opt.lr, beta1=opt.b1, beta2=opt.b2)

G_Optim.update_parameters_name('G_Optim')
D_Optim.update_parameters_name('D_Optim')


def compute_gradient_penalty(real_samples, fake_samples):
    """Calculates the gradient penalty loss for WGAN GP"""
    # Random weight term for interpolation between real and fake samples
    alpha = ops.randn((real_samples.shape[0], 1, 1, 1))
    # Get random interpolation between real and fake samples
    interpolates = alpha * real_samples + ((1 - alpha) * fake_samples)

    grad_fn = ops.grad(discriminator)
    # Get gradient w.r.t. interpolates
    gradients = grad_fn(interpolates)

    gradients = gradients.view(gradients.shape[0], -1)
    gradient_penalty = ((mindspore.numpy.norm(gradients, 2, axis=1) - 1) ** 2).mean()
    return gradient_penalty


def g_forward(_z):
    """Generator forward function"""
    _fake_imgs = generator(_z)
    # Loss measures generator's ability to fool the discriminator
    # Train on fake images
    fake_validity = discriminator(_fake_imgs)
    _g_loss = -ops.mean(fake_validity)
    return _g_loss, _fake_imgs


# 判别器正向传播
def d_forward(_real_imgs):
    """Discriminator forward function"""
    _z = ops.StandardNormal()((imgs.shape[0], opt.latent_dim))
    fake_imgs = generator(_z)
    # Real images
    real_validity = discriminator(_real_imgs)
    # Fake images
    fake_validity = discriminator(fake_imgs)
    # Gradient penalty
    gradient_penalty = compute_gradient_penalty(_real_imgs, fake_imgs)
    # Adversarial loss
    _d_loss = -ops.mean(real_validity) + ops.mean(fake_validity) + lambda_gp * gradient_penalty

    return _d_loss, _z


transform = [
    transforms.Rescale(1.0 / 255.0, 0),
    transforms.Resize(opt.img_size),
    transforms.Normalize([0.5], [0.5]),
    transforms.HWC2CHW()
]

dataset = mindspore.dataset.MnistDataset(
    dataset_dir=file_path,
    usage='train',
    shuffle=True
).map(operations=transform, input_columns="image").batch(opt.batch_size)

grad_g = ops.value_and_grad(g_forward, None, G_Optim.parameters, has_aux=True)
grad_d = ops.value_and_grad(d_forward, None, D_Optim.parameters, has_aux=True)

batches_done = 0

generator.update_parameters_name('generator')
discriminator.update_parameters_name('discriminator')
discriminator.set_train()
generator.set_train()

for epoch in range(opt.n_epochs):
    for i, (imgs, _) in enumerate(dataset.create_tuple_iterator()):
        real_imgs = Tensor(imgs, dtype=mstype.float32)

        (loss_D, z), D_grads = grad_d(real_imgs)
        D_Optim(D_grads)

        # Train the generator every n_critic iterations
        if i % opt.n_critic == 0:
            (loss_G, gen_imgs), G_grads = grad_g(z)
            G_Optim(G_grads)

            print(
                f'[Epoch {epoch}/{opt.n_epochs}] [Batch {i}/{dataset.get_dataset_size()}] '
                f'[D loss: {loss_D.asnumpy().item()}] [G loss: {loss_G.asnumpy().item()}]'
            )

            if batches_done % opt.sample_interval == 0:
                save_imgs(gen_imgs[:25].asnumpy(), batches_done)
                # to_image(gen_imgs[:25], os.path.join("images", F'{batches_done}.png'))
            batches_done += opt.n_critic

implementations/wgan_gp/wgan_gp.py
mindspore-courses/MindSpore-GAN