<gh_stars>10-100
import torch.nn as nn
import torch
import random

from typing import List


def convert_bn_to_in(module: nn.Module):
    """
    Convert batch_norm_2d in a given module to instance_norm_2d
    :param module:
    :return:
    """
    module_output = module
    if isinstance(module, nn.BatchNorm2d):
        module_output = nn.InstanceNorm2d(module.num_features, module.eps, affine=module.affine,
                                          track_running_stats=module.track_running_stats)
    for name, child in module.named_children():
        module_output.add_module(name, convert_bn_to_in(child))
    del module
    return module_output


def _set_requires_grad(models: [List[nn.Module], nn.Module], requires_grad: bool):
    """
    Freeze model's parameters
    :param models: model or list of models
    :param requires_grad: if requires grad
    :return:
    """
    if not isinstance(models, (list, tuple)):
        models = [models]
    for m in models:
        if m is not None:
            for param in m.parameters():
                param.requires_grad = requires_grad


def freeze_model(models):
    _set_requires_grad(models, False)


def unfreeze_model(models):
    _set_requires_grad(models, True)


def get_norm_layer(norm_type: str, num_features: int):
    if norm_type == 'batch_norm':
        return nn.BatchNorm2d(num_features, affine=True)
    elif norm_type == 'instance_norm':
        return nn.InstanceNorm2d(num_features, affine=True)
    else:
        raise ValueError('Unsupported norm layer type: {}.'.format(norm_type))


def get_non_linear(non_linear: str):
    if non_linear == 'elu':
        return nn.ELU(inplace=True)
    elif non_linear == 'relu':
        return nn.ReLU(inplace=True)
    elif non_linear == 'leaky_relu':
        return nn.LeakyReLU(inplace=True)
    else:
        raise ValueError('Unsupported non_linear layer: {}.'.format(non_linear))


class ImagePool:
    """This class implements an image buffer that stores previously generated images.

    This buffer enables us to update discriminators using a history of generated images
    rather than the ones produced by the latest generators.
    """

    def __init__(self, pool_size):
        """Initialize the ImagePool class

        Parameters:
            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
        """
        self.pool_size = pool_size
        if self.pool_size > 0:  # create an empty pool
            self.num_imgs = 0
            self.images = []

    def query(self, images):
        """Return an image from the pool.

        Parameters:
            images: the latest generated images from the generator

        Returns images from the buffer.

        By 50/100, the buffer will return input images.
        By 50/100, the buffer will return images previously stored in the buffer,
        and insert the current images to the buffer.
        """
        if self.pool_size == 0:  # if the buffer size is 0, do nothing
            return images
        return_images = []
        for image in images:
            image = torch.unsqueeze(image.data, 0)
            if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer
                self.num_imgs = self.num_imgs + 1
                self.images.append(image)
                return_images.append(image)
            else:
                p = random.uniform(0, 1)
                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
                    tmp = self.images[random_id].clone()
                    self.images[random_id] = image
                    return_images.append(tmp)
                else:       # by another 50% chance, the buffer will return the current image
                    return_images.append(image)
        return_images = torch.cat(return_images, 0)   # collect all the images and return
        return return_images


import torch.nn as nn
import torch
import random

from typing import List


def convert_bn_to_in(module: nn.Module):
    """
    Convert batch_norm_2d in a given module to instance_norm_2d
    :param module:
    :return:
    """
    module_output = module
    if isinstance(module, nn.BatchNorm2d):
        module_output = nn.InstanceNorm2d(module.num_features, module.eps, affine=module.affine,
                                          track_running_stats=module.track_running_stats)
    for name, child in module.named_children():
        module_output.add_module(name, convert_bn_to_in(child))
    del module
    return module_output


def _set_requires_grad(models: [List[nn.Module], nn.Module], requires_grad: bool):
    """
    Freeze model's parameters
    :param models: model or list of models
    :param requires_grad: if requires grad
    :return:
    """
    if not isinstance(models, (list, tuple)):
        models = [models]
    for m in models:
        if m is not None:
            for param in m.parameters():
                param.requires_grad = requires_grad


def freeze_model(models):
    _set_requires_grad(models, False)


def unfreeze_model(models):
    _set_requires_grad(models, True)


def get_norm_layer(norm_type: str, num_features: int):
    if norm_type == 'batch_norm':
        return nn.BatchNorm2d(num_features, affine=True)
    elif norm_type == 'instance_norm':
        return nn.InstanceNorm2d(num_features, affine=True)
    else:
        raise ValueError('Unsupported norm layer type: {}.'.format(norm_type))


def get_non_linear(non_linear: str):
    if non_linear == 'elu':
        return nn.ELU(inplace=True)
    elif non_linear == 'relu':
        return nn.ReLU(inplace=True)
    elif non_linear == 'leaky_relu':
        return nn.LeakyReLU(inplace=True)
    else:
        raise ValueError('Unsupported non_linear layer: {}.'.format(non_linear))


class ImagePool:
    """This class implements an image buffer that stores previously generated images.

    This buffer enables us to update discriminators using a history of generated images
    rather than the ones produced by the latest generators.
    """

    def __init__(self, pool_size):
        """Initialize the ImagePool class

        Parameters:
            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
        """
        self.pool_size = pool_size
        if self.pool_size > 0:  # create an empty pool
            self.num_imgs = 0
            self.images = []

    def query(self, images):
        """Return an image from the pool.

        Parameters:
            images: the latest generated images from the generator

        Returns images from the buffer.

        By 50/100, the buffer will return input images.
        By 50/100, the buffer will return images previously stored in the buffer,
        and insert the current images to the buffer.
        """
        if self.pool_size == 0:  # if the buffer size is 0, do nothing
            return images
        return_images = []
        for image in images:
            image = torch.unsqueeze(image.data, 0)
            if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer
                self.num_imgs = self.num_imgs + 1
                self.images.append(image)
                return_images.append(image)
            else:
                p = random.uniform(0, 1)
                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
                    tmp = self.images[random_id].clone()
                    self.images[random_id] = image
                    return_images.append(tmp)
                else:       # by another 50% chance, the buffer will return the current image
                    return_images.append(image)
        return_images = torch.cat(return_images, 0)   # collect all the images and return
        return return_images
Convert batch_norm_2d in a given module to instance_norm_2d     :param module:     :return: Freeze model's parameters     :param models: model or list of models     :param requires_grad: if requires grad     :return: This class implements an image buffer that stores previously generated images.      This buffer enables us to update discriminators using a history of generated images     rather than the ones produced by the latest generators. Initialize the ImagePool class          Parameters:             pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created # create an empty pool Return an image from the pool.          Parameters:             images: the latest generated images from the generator          Returns images from the buffer.          By 50/100, the buffer will return input images.         By 50/100, the buffer will return images previously stored in the buffer,         and insert the current images to the buffer. # if the buffer size is 0, do nothing # if the buffer is not full; keep inserting current images to the buffer # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer # randint is inclusive # by another 50% chance, the buffer will return the current image # collect all the images and return