Myosotis-Researches 0.1.7__tar.gz → 0.1.9__tar.gz

{myosotis_researches-0.1.7 → myosotis_researches-0.1.9/Myosotis_Researches.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Myosotis-Researches
-Version: 0.1.7
+Version: 0.1.9
 Summary: A repository for storing my progress of researches.
 Home-page: https://github.com/Zeyu-Xie/Myosotis-Researches
 Author: Zeyu Xie

myosotis_researches-0.1.9/Myosotis_Researches.egg-info/SOURCES.txt

@@ -0,0 +1,27 @@
+LICENSE
+README.md
+pyproject.toml
+setup.py
+Myosotis_Researches.egg-info/PKG-INFO
+Myosotis_Researches.egg-info/SOURCES.txt
+Myosotis_Researches.egg-info/dependency_links.txt
+Myosotis_Researches.egg-info/top_level.txt
+myosotis_researches/__init__.py
+myosotis_researches/CcGAN/__init__.py
+myosotis_researches/CcGAN/internal/__init__.py
+myosotis_researches/CcGAN/internal/install_datasets.py
+myosotis_researches/CcGAN/internal/show_datasets.py
+myosotis_researches/CcGAN/internal/uninstall_datasets.py
+myosotis_researches/CcGAN/train/__init__.py
+myosotis_researches/CcGAN/train/train_ccgan.py
+myosotis_researches/CcGAN/train/train_cgan.py
+myosotis_researches/CcGAN/train/train_cgan_concat.py
+myosotis_researches/CcGAN/train/train_net_for_label_embed.py
+myosotis_researches/CcGAN/utils/IMGs_dataset.py
+myosotis_researches/CcGAN/utils/SimpleProgressBar.py
+myosotis_researches/CcGAN/utils/__init__.py
+myosotis_researches/CcGAN/utils/concat_image.py
+myosotis_researches/CcGAN/utils/make_h5.py
+myosotis_researches/CcGAN/utils/opts.py
+myosotis_researches/CcGAN/utils/print_hdf5.py
+myosotis_researches/CcGAN/utils/train.py

{myosotis_researches-0.1.7/Myosotis_Researches.egg-info → myosotis_researches-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Myosotis-Researches
-Version: 0.1.7
+Version: 0.1.9
 Summary: A repository for storing my progress of researches.
 Home-page: https://github.com/Zeyu-Xie/Myosotis-Researches
 Author: Zeyu Xie

myosotis_researches-0.1.9/myosotis_researches/CcGAN/train/__init__.py

@@ -0,0 +1,4 @@
+from .train_ccgan import *
+from .train_cgan import *
+from .train_cgan_concat import *
+from .train_net_for_label_embed import *

{myosotis_researches-0.1.7/myosotis_researches/CcGAN/train_128_output_10 → myosotis_researches-0.1.9/myosotis_researches/CcGAN/train}/train_ccgan.py

@@ -5,10 +5,9 @@ import timeit
 from PIL import Image
 from torchvision.utils import save_image
 import torch.cuda as cutorch
+import sys
 
-from .utils import SimpleProgressBar, IMGs_dataset
-from .opts import parse_opts
-from .DiffAugment_pytorch import DiffAugment
+from myosotis_researches.CcGAN.utils import *
 
 ''' Settings '''
 args = parse_opts()
@@ -79,7 +78,8 @@ def train_ccgan(kernel_sigma, kappa, train_images, train_labels, netG, netD, net
     # printed images with labels between the 5-th quantile and 95-th quantile of training labels
     n_row=10; n_col = 1
     z_fixed = torch.randn(n_row*n_col, dim_gan, dtype=torch.float).cuda()
-    selected_labels = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+
+    selected_labels = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
 
     y_fixed = np.zeros(n_row*n_col)
     for i in range(n_row):

{myosotis_researches-0.1.7/myosotis_researches/CcGAN/train_128 → myosotis_researches-0.1.9/myosotis_researches/CcGAN/train}/train_cgan.py

@@ -6,9 +6,7 @@ import numpy as np
 import os
 import timeit
 
-from .utils import IMGs_dataset, SimpleProgressBar
-from .opts import parse_opts
-from .DiffAugment_pytorch import DiffAugment
+from myosotis_researches.CcGAN.utils import *
 
 ''' Settings '''
 args = parse_opts()

{myosotis_researches-0.1.7/myosotis_researches/CcGAN/train_128 → myosotis_researches-0.1.9/myosotis_researches/CcGAN/train}/train_cgan_concat.py

@@ -6,9 +6,7 @@ import numpy as np
 import os
 import timeit
 
-from .utils import IMGs_dataset, SimpleProgressBar
-from .opts import parse_opts
-from .DiffAugment_pytorch import DiffAugment
+from myosotis_researches.CcGAN.utils import *
 
 ''' Settings '''
 args = parse_opts()

{myosotis_researches-0.1.7 → myosotis_researches-0.1.9}/myosotis_researches/CcGAN/utils/__init__.py

@@ -3,7 +3,7 @@ from .concat_image import concat_image
 from .make_h5 import make_h5
 from .SimpleProgressBar import SimpleProgressBar
 from .IMGs_dataset import IMGs_dataset
-from .train import PlotLoss, compute_entropy, predict_class_labels
+from .train import PlotLoss, compute_entropy, predict_class_labels, DiffAugment
 from .opts import parse_opts
 
 __all__ = [
@@ -15,5 +15,6 @@ __all__ = [
     "PlotLoss",
     "compute_entropy",
     "predict_class_labels",
+    "DiffAugment",
     "parse_opts"
 ]

myosotis_researches-0.1.9/myosotis_researches/CcGAN/utils/train.py

@@ -0,0 +1,156 @@
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def PlotLoss(loss, filename):
+    x_axis = np.arange(start=1, stop=len(loss) + 1)
+    plt.switch_backend("agg")
+    mpl.style.use("seaborn")
+    fig = plt.figure()
+    ax = plt.subplot(111)
+    ax.plot(x_axis, np.array(loss))
+    plt.xlabel("epoch")
+    plt.ylabel("training loss")
+    plt.legend()
+    # ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15),  shadow=True, ncol=3)
+    # plt.title('Training Loss')
+    plt.savefig(filename)
+
+
+# compute entropy of class labels; labels is a numpy array
+def compute_entropy(labels, base=None):
+    value, counts = np.unique(labels, return_counts=True)
+    norm_counts = counts / counts.sum()
+    base = np.e if base is None else base
+    return -(norm_counts * np.log(norm_counts) / np.log(base)).sum()
+
+
+def predict_class_labels(net, images, batch_size=500, verbose=False, num_workers=0):
+    net = net.cuda()
+    net.eval()
+
+    n = len(images)
+    if batch_size > n:
+        batch_size = n
+    dataset_pred = IMGs_dataset(images, normalize=False)
+    dataloader_pred = torch.utils.data.DataLoader(
+        dataset_pred, batch_size=batch_size, shuffle=False, num_workers=num_workers
+    )
+
+    class_labels_pred = np.zeros(n + batch_size)
+    with torch.no_grad():
+        nimgs_got = 0
+        if verbose:
+            pb = SimpleProgressBar()
+        for batch_idx, batch_images in enumerate(dataloader_pred):
+            batch_images = batch_images.type(torch.float).cuda()
+            batch_size_curr = len(batch_images)
+
+            outputs, _ = net(batch_images)
+            _, batch_class_labels_pred = torch.max(outputs.data, 1)
+            class_labels_pred[nimgs_got : (nimgs_got + batch_size_curr)] = (
+                batch_class_labels_pred.detach().cpu().numpy().reshape(-1)
+            )
+
+            nimgs_got += batch_size_curr
+            if verbose:
+                pb.update((float(nimgs_got) / n) * 100)
+        # end for batch_idx
+    class_labels_pred = class_labels_pred[0:n]
+    return class_labels_pred
+
+
+def DiffAugment(x, policy="", channels_first=True):
+    if policy:
+        if not channels_first:
+            x = x.permute(0, 3, 1, 2)
+        for p in policy.split(","):
+            for f in AUGMENT_FNS[p]:
+                x = f(x)
+        if not channels_first:
+            x = x.permute(0, 2, 3, 1)
+        x = x.contiguous()
+    return x
+
+
+def rand_brightness(x):
+    x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
+    return x
+
+
+def rand_saturation(x):
+    x_mean = x.mean(dim=1, keepdim=True)
+    x = (x - x_mean) * (
+        torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2
+    ) + x_mean
+    return x
+
+
+def rand_contrast(x):
+    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
+    x = (x - x_mean) * (
+        torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5
+    ) + x_mean
+    return x
+
+
+def rand_translation(x, ratio=0.125):
+    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    translation_x = torch.randint(
+        -shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device
+    )
+    translation_y = torch.randint(
+        -shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device
+    )
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(x.size(2), dtype=torch.long, device=x.device),
+        torch.arange(x.size(3), dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
+    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
+    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
+    x = (
+        x_pad.permute(0, 2, 3, 1)
+        .contiguous()[grid_batch, grid_x, grid_y]
+        .permute(0, 3, 1, 2)
+    )
+    return x
+
+
+def rand_cutout(x, ratio=0.5):
+    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    offset_x = torch.randint(
+        0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device
+    )
+    offset_y = torch.randint(
+        0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device
+    )
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(
+        grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1
+    )
+    grid_y = torch.clamp(
+        grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1
+    )
+    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
+    mask[grid_batch, grid_x, grid_y] = 0
+    x = x * mask.unsqueeze(1)
+    return x
+
+
+AUGMENT_FNS = {
+    "color": [rand_brightness, rand_saturation, rand_contrast],
+    "translation": [rand_translation],
+    "cutout": [rand_cutout],
+}
+
+
+__all__ = ["PlotLoss", "compute_entropy", "predict_class_labels", "DiffAugment"]

{myosotis_researches-0.1.7 → myosotis_researches-0.1.9}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="Myosotis-Researches",
-    version="0.1.7",
+    version="0.1.9",
     description="A repository for storing my progress of researches.",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",

myosotis_researches-0.1.7/Myosotis_Researches.egg-info/SOURCES.txt

@@ -1,62 +0,0 @@
-LICENSE
-README.md
-pyproject.toml
-setup.py
-Myosotis_Researches.egg-info/PKG-INFO
-Myosotis_Researches.egg-info/SOURCES.txt
-Myosotis_Researches.egg-info/dependency_links.txt
-Myosotis_Researches.egg-info/top_level.txt
-myosotis_researches/__init__.py
-myosotis_researches/CcGAN/__init__.py
-myosotis_researches/CcGAN/internal/__init__.py
-myosotis_researches/CcGAN/internal/install_datasets.py
-myosotis_researches/CcGAN/internal/show_datasets.py
-myosotis_researches/CcGAN/internal/uninstall_datasets.py
-myosotis_researches/CcGAN/models_128/CcGAN_SAGAN.py
-myosotis_researches/CcGAN/models_128/ResNet_class_eval.py
-myosotis_researches/CcGAN/models_128/ResNet_embed.py
-myosotis_researches/CcGAN/models_128/ResNet_regre_eval.py
-myosotis_researches/CcGAN/models_128/__init__.py
-myosotis_researches/CcGAN/models_128/autoencoder.py
-myosotis_researches/CcGAN/models_128/cGAN_SAGAN.py
-myosotis_researches/CcGAN/models_128/cGAN_concat_SAGAN.py
-myosotis_researches/CcGAN/models_256/CcGAN_SAGAN.py
-myosotis_researches/CcGAN/models_256/ResNet_class_eval.py
-myosotis_researches/CcGAN/models_256/ResNet_embed.py
-myosotis_researches/CcGAN/models_256/ResNet_regre_eval.py
-myosotis_researches/CcGAN/models_256/__init__.py
-myosotis_researches/CcGAN/models_256/autoencoder.py
-myosotis_researches/CcGAN/models_256/cGAN_SAGAN.py
-myosotis_researches/CcGAN/models_256/cGAN_concat_SAGAN.py
-myosotis_researches/CcGAN/train_128/DiffAugment_pytorch.py
-myosotis_researches/CcGAN/train_128/__init__.py
-myosotis_researches/CcGAN/train_128/eval_metrics.py
-myosotis_researches/CcGAN/train_128/opts.py
-myosotis_researches/CcGAN/train_128/pretrain_AE.py
-myosotis_researches/CcGAN/train_128/pretrain_CNN_class.py
-myosotis_researches/CcGAN/train_128/pretrain_CNN_regre.py
-myosotis_researches/CcGAN/train_128/train_ccgan.py
-myosotis_researches/CcGAN/train_128/train_cgan.py
-myosotis_researches/CcGAN/train_128/train_cgan_concat.py
-myosotis_researches/CcGAN/train_128/train_net_for_label_embed.py
-myosotis_researches/CcGAN/train_128/utils.py
-myosotis_researches/CcGAN/train_128_output_10/DiffAugment_pytorch.py
-myosotis_researches/CcGAN/train_128_output_10/__init__.py
-myosotis_researches/CcGAN/train_128_output_10/eval_metrics.py
-myosotis_researches/CcGAN/train_128_output_10/opts.py
-myosotis_researches/CcGAN/train_128_output_10/pretrain_AE.py
-myosotis_researches/CcGAN/train_128_output_10/pretrain_CNN_class.py
-myosotis_researches/CcGAN/train_128_output_10/pretrain_CNN_regre.py
-myosotis_researches/CcGAN/train_128_output_10/train_ccgan.py
-myosotis_researches/CcGAN/train_128_output_10/train_cgan.py
-myosotis_researches/CcGAN/train_128_output_10/train_cgan_concat.py
-myosotis_researches/CcGAN/train_128_output_10/train_net_for_label_embed.py
-myosotis_researches/CcGAN/train_128_output_10/utils.py
-myosotis_researches/CcGAN/utils/IMGs_dataset.py
-myosotis_researches/CcGAN/utils/SimpleProgressBar.py
-myosotis_researches/CcGAN/utils/__init__.py
-myosotis_researches/CcGAN/utils/concat_image.py
-myosotis_researches/CcGAN/utils/make_h5.py
-myosotis_researches/CcGAN/utils/opts.py
-myosotis_researches/CcGAN/utils/print_hdf5.py
-myosotis_researches/CcGAN/utils/train.py

myosotis_researches-0.1.7/myosotis_researches/CcGAN/models_128/CcGAN_SAGAN.py

@@ -1,301 +0,0 @@
-'''
-
-Adapted from https://github.com/voletiv/self-attention-GAN-pytorch/blob/master/sagan_models.py
-
-
-'''
-
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from torch.nn.utils import spectral_norm
-from torch.nn.init import xavier_uniform_
-
-
-def init_weights(m):
-    if type(m) == nn.Linear or type(m) == nn.Conv2d:
-        xavier_uniform_(m.weight)
-        if m.bias is not None:
-            m.bias.data.fill_(0.)
-
-
-def snconv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
-    return spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
-                                   stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias))
-
-def snlinear(in_features, out_features, bias=True):
-    return spectral_norm(nn.Linear(in_features=in_features, out_features=out_features, bias=bias))
-
-
-
-class Self_Attn(nn.Module):
-    """ Self attention Layer"""
-
-    def __init__(self, in_channels):
-        super(Self_Attn, self).__init__()
-        self.in_channels = in_channels
-        self.snconv1x1_theta = snconv2d(in_channels=in_channels, out_channels=in_channels//8, kernel_size=1, stride=1, padding=0)
-        self.snconv1x1_phi = snconv2d(in_channels=in_channels, out_channels=in_channels//8, kernel_size=1, stride=1, padding=0)
-        self.snconv1x1_g = snconv2d(in_channels=in_channels, out_channels=in_channels//2, kernel_size=1, stride=1, padding=0)
-        self.snconv1x1_attn = snconv2d(in_channels=in_channels//2, out_channels=in_channels, kernel_size=1, stride=1, padding=0)
-        self.maxpool = nn.MaxPool2d(2, stride=2, padding=0)
-        self.softmax  = nn.Softmax(dim=-1)
-        self.sigma = nn.Parameter(torch.zeros(1))
-
-    def forward(self, x):
-        """
-            inputs :
-                x : input feature maps(B X C X W X H)
-            returns :
-                out : self attention value + input feature
-                attention: B X N X N (N is Width*Height)
-        """
-        _, ch, h, w = x.size()
-        # Theta path
-        theta = self.snconv1x1_theta(x)
-        theta = theta.view(-1, ch//8, h*w)
-        # Phi path
-        phi = self.snconv1x1_phi(x)
-        phi = self.maxpool(phi)
-        phi = phi.view(-1, ch//8, h*w//4)
-        # Attn map
-        attn = torch.bmm(theta.permute(0, 2, 1), phi)
-        attn = self.softmax(attn)
-        # g path
-        g = self.snconv1x1_g(x)
-        g = self.maxpool(g)
-        g = g.view(-1, ch//2, h*w//4)
-        # Attn_g
-        attn_g = torch.bmm(g, attn.permute(0, 2, 1))
-        attn_g = attn_g.view(-1, ch//2, h, w)
-        attn_g = self.snconv1x1_attn(attn_g)
-        # Out
-        out = x + self.sigma*attn_g
-        return out
-
-
-
-
-'''
-
-Generator
-
-'''
-
-
-class ConditionalBatchNorm2d(nn.Module):
-    def __init__(self, num_features, dim_embed):
-        super().__init__()
-        self.num_features = num_features
-        self.bn = nn.BatchNorm2d(num_features, momentum=0.001, affine=False)
-        self.embed_gamma = nn.Linear(dim_embed, num_features, bias=False)
-        self.embed_beta = nn.Linear(dim_embed, num_features, bias=False)
-
-    def forward(self, x, y):
-        out = self.bn(x)
-        gamma = self.embed_gamma(y).view(-1, self.num_features, 1, 1)
-        beta = self.embed_beta(y).view(-1, self.num_features, 1, 1)
-        out = out + gamma*out + beta
-        return out
-
-
-class GenBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, dim_embed):
-        super(GenBlock, self).__init__()
-        self.cond_bn1 = ConditionalBatchNorm2d(in_channels, dim_embed)
-        self.relu = nn.ReLU(inplace=True)
-        self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
-        self.cond_bn2 = ConditionalBatchNorm2d(out_channels, dim_embed)
-        self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
-        self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x, labels):
-        x0 = x
-
-        x = self.cond_bn1(x, labels)
-        x = self.relu(x)
-        x = F.interpolate(x, scale_factor=2, mode='nearest') # upsample
-        x = self.snconv2d1(x)
-        x = self.cond_bn2(x, labels)
-        x = self.relu(x)
-        x = self.snconv2d2(x)
-
-        x0 = F.interpolate(x0, scale_factor=2, mode='nearest') # upsample
-        x0 = self.snconv2d0(x0)
-
-        out = x + x0
-        return out
-
-
-class CcGAN_SAGAN_Generator(nn.Module):
-    """Generator."""
-
-    def __init__(self, dim_z, dim_embed=128, nc=3, gene_ch=64):
-        super(CcGAN_SAGAN_Generator, self).__init__()
-
-        self.dim_z = dim_z
-        self.gene_ch = gene_ch
-
-        self.snlinear0 = snlinear(in_features=dim_z, out_features=gene_ch*16*4*4)
-        self.block1 = GenBlock(gene_ch*16, gene_ch*16, dim_embed)
-        self.block2 = GenBlock(gene_ch*16, gene_ch*8, dim_embed)
-        self.block3 = GenBlock(gene_ch*8, gene_ch*4, dim_embed)
-        self.self_attn = Self_Attn(gene_ch*4)
-        self.block4 = GenBlock(gene_ch*4, gene_ch*2, dim_embed)
-        self.block5 = GenBlock(gene_ch*2, gene_ch, dim_embed)
-        self.bn = nn.BatchNorm2d(gene_ch, eps=1e-5, momentum=0.0001, affine=True)
-        self.relu = nn.ReLU(inplace=True)
-        self.snconv2d1 = snconv2d(in_channels=gene_ch, out_channels=nc, kernel_size=3, stride=1, padding=1)
-        self.tanh = nn.Tanh()
-
-        # Weight init
-        self.apply(init_weights)
-
-    def forward(self, z, labels):
-        # n x dim_z
-        out = self.snlinear0(z)            # 4*4
-        out = out.view(-1, self.gene_ch*16, 4, 4) # 4 x 4
-        out = self.block1(out, labels)    # 8 x 8
-        out = self.block2(out, labels)    # 16 x 16
-        out = self.block3(out, labels)    # 32 x 32
-        out = self.self_attn(out)         # 32 x 32
-        out = self.block4(out, labels)    # 64 x 64
-        out = self.block5(out, labels)    # 128 x 128
-        out = self.bn(out)
-        out = self.relu(out)
-        out = self.snconv2d1(out)
-        out = self.tanh(out)
-        return out
-
-
-
-'''
-
-Discriminator
-
-'''
-
-class DiscOptBlock(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(DiscOptBlock, self).__init__()
-        self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
-        self.relu = nn.ReLU(inplace=True)
-        self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
-        self.downsample = nn.AvgPool2d(2)
-        self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x):
-        x0 = x
-
-        x = self.snconv2d1(x)
-        x = self.relu(x)
-        x = self.snconv2d2(x)
-        x = self.downsample(x)
-
-        x0 = self.downsample(x0)
-        x0 = self.snconv2d0(x0)
-
-        out = x + x0
-        return out
-
-
-class DiscBlock(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(DiscBlock, self).__init__()
-        self.relu = nn.ReLU(inplace=True)
-        self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
-        self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
-        self.downsample = nn.AvgPool2d(2)
-        self.ch_mismatch = False
-        if in_channels != out_channels:
-            self.ch_mismatch = True
-        self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x, downsample=True):
-        x0 = x
-
-        x = self.relu(x)
-        x = self.snconv2d1(x)
-        x = self.relu(x)
-        x = self.snconv2d2(x)
-        if downsample:
-            x = self.downsample(x)
-
-        if downsample or self.ch_mismatch:
-            x0 = self.snconv2d0(x0)
-            if downsample:
-                x0 = self.downsample(x0)
-
-        out = x + x0
-        return out
-
-
-class CcGAN_SAGAN_Discriminator(nn.Module):
-    """Discriminator."""
-
-    def __init__(self, dim_embed=128, nc=3, disc_ch=64):
-        super(CcGAN_SAGAN_Discriminator, self).__init__()
-        self.disc_ch = disc_ch
-        self.opt_block1 = DiscOptBlock(nc, disc_ch)
-        self.block1 = DiscBlock(disc_ch, disc_ch*2)
-        self.self_attn = Self_Attn(disc_ch*2)
-        self.block2 = DiscBlock(disc_ch*2, disc_ch*4)
-        self.block3 = DiscBlock(disc_ch*4, disc_ch*8)
-        self.block4 = DiscBlock(disc_ch*8, disc_ch*16)
-        self.block5 = DiscBlock(disc_ch*16, disc_ch*16)
-        self.relu = nn.ReLU(inplace=True)
-        self.snlinear1 = snlinear(in_features=disc_ch*16*4*4, out_features=1)
-        self.sn_embedding1 = snlinear(dim_embed, disc_ch*16*4*4, bias=False)
-
-        # Weight init
-        self.apply(init_weights)
-        xavier_uniform_(self.sn_embedding1.weight)
-
-    def forward(self, x, labels):
-        # 128x128
-        out = self.opt_block1(x)  # 128x128
-        out = self.block1(out)    # 64 x 64
-        out = self.self_attn(out) # 64 x 64
-        out = self.block2(out)    # 32 x 32
-        out = self.block3(out)    # 16 x 16
-        out = self.block4(out)    # 8 x 8
-        out = self.block5(out, downsample=False)    # 4 x 4
-        out = self.relu(out)              # n x disc_ch*16 x 4 x 4
-        out = out.view(-1, self.disc_ch*16*4*4)
-        output1 = torch.squeeze(self.snlinear1(out)) # n
-        # Projection
-        h_labels = self.sn_embedding1(labels)   # n x disc_ch*16
-        proj = torch.mul(out, h_labels)          # n x disc_ch*16
-        output2 = torch.sum(proj, dim=[1])      # n
-        # Out
-        output = output1 + output2              # n
-        return output
-
-
-if __name__ == "__main__":
-
-    def get_parameter_number(net):
-        total_num = sum(p.numel() for p in net.parameters())
-        trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
-        return {'Total': total_num, 'Trainable': trainable_num}
-
-
-    netG = CcGAN_SAGAN_Generator(dim_z=256, dim_embed=128, gene_ch=128).cuda()
-    netD = CcGAN_SAGAN_Discriminator(dim_embed=128, disc_ch=128).cuda()
-
-    # netG = nn.DataParallel(netG)
-    # netD = nn.DataParallel(netD)
-
-    N=4
-    z = torch.randn(N, 256).cuda()
-    y = torch.randn(N, 128).cuda()
-    x = netG(z,y)
-    o = netD(x,y)
-    print(x.size())
-    print(o.size())
-
-    print('G:', get_parameter_number(netG))
-    print('D:', get_parameter_number(netD))