Myosotis-Researches 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

myosotis_researches/CcGAN/train_128/train_cgan.py

@@ -1,254 +0,0 @@
-
-import torch
-import torch.nn as nn
-from torchvision.utils import save_image
-import numpy as np
-import os
-import timeit
-
-from .utils import IMGs_dataset, SimpleProgressBar
-from .opts import parse_opts
-from .DiffAugment_pytorch import DiffAugment
-
-''' Settings '''
-args = parse_opts()
-
-# some parameters in opts
-loss_type = args.loss_type_gan
-niters = args.niters_gan
-resume_niters = args.resume_niters_gan
-dim_gan = args.dim_gan
-lr_g = args.lr_g_gan
-lr_d = args.lr_d_gan
-save_niters_freq = args.save_niters_freq
-batch_size = min(args.batch_size_disc, args.batch_size_gene)
-num_classes = args.cGAN_num_classes
-gan_arch = args.GAN_arch
-num_D_steps = args.num_D_steps
-
-visualize_freq = args.visualize_freq
-
-num_workers = args.num_workers
-
-NC = args.num_channels
-IMG_SIZE = args.img_size
-
-use_DiffAugment = args.gan_DiffAugment
-policy = args.gan_DiffAugment_policy
-
-
-def train_cgan(images, labels, netG, netD, save_images_folder, save_models_folder = None):
-
-    netG = netG.cuda()
-    netD = netD.cuda()
-
-    optimizerG = torch.optim.Adam(netG.parameters(), lr=lr_g, betas=(0.5, 0.999))
-    optimizerD = torch.optim.Adam(netD.parameters(), lr=lr_d, betas=(0.5, 0.999))
-
-    trainset = IMGs_dataset(images, labels, normalize=True)
-    train_dataloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
-    unique_labels = np.sort(np.array(list(set(labels)))).astype(np.int)
-
-    if save_models_folder is not None and resume_niters>0:
-        save_file = save_models_folder + "/cGAN_{}_nDsteps_{}_checkpoint_intrain/cGAN_checkpoint_niters_{}.pth".format(gan_arch, num_D_steps, resume_niters)
-        checkpoint = torch.load(save_file)
-        netG.load_state_dict(checkpoint['netG_state_dict'])
-        netD.load_state_dict(checkpoint['netD_state_dict'])
-        optimizerG.load_state_dict(checkpoint['optimizerG_state_dict'])
-        optimizerD.load_state_dict(checkpoint['optimizerD_state_dict'])
-        torch.set_rng_state(checkpoint['rng_state'])
-    #end if
-
-    n_row=10
-    z_fixed = torch.randn(n_row**2, dim_gan, dtype=torch.float).cuda()
-    unique_labels = np.sort(unique_labels)
-    selected_labels = np.zeros(n_row)
-    indx_step_size = len(unique_labels)//n_row
-    for i in range(n_row):
-        indx = i*indx_step_size
-        selected_labels[i] = unique_labels[indx]
-    y_fixed = np.zeros(n_row**2)
-    for i in range(n_row):
-        curr_label = selected_labels[i]
-        for j in range(n_row):
-            y_fixed[i*n_row+j] = curr_label
-    y_fixed = torch.from_numpy(y_fixed).type(torch.long).cuda()
-
-
-    batch_idx = 0
-    dataloader_iter = iter(train_dataloader)
-
-    start_time = timeit.default_timer()
-    for niter in range(resume_niters, niters):
-
-        if batch_idx+1 == len(train_dataloader):
-            dataloader_iter = iter(train_dataloader)
-            batch_idx = 0
-
-        '''
-
-        Train Generator: maximize log(D(G(z)))
-
-        '''
-
-        netG.train()
-
-        # get training images
-        _, batch_train_labels = dataloader_iter.next()
-        assert batch_size == batch_train_labels.shape[0]
-        batch_train_labels = batch_train_labels.type(torch.long).cuda()
-        batch_idx+=1
-
-        # Sample noise and labels as generator input
-        z = torch.randn(batch_size, dim_gan, dtype=torch.float).cuda()
-
-        #generate fake images
-        batch_fake_images = netG(z, batch_train_labels)
-
-        # Loss measures generator's ability to fool the discriminator
-        if use_DiffAugment:
-            dis_out = netD(DiffAugment(batch_fake_images, policy=policy), batch_train_labels)
-        else:
-            dis_out = netD(batch_fake_images, batch_train_labels)
-        
-        if loss_type == "vanilla":
-            dis_out = torch.nn.Sigmoid()(dis_out)
-            g_loss = - torch.mean(torch.log(dis_out+1e-20))
-        elif loss_type == "hinge":
-            g_loss = - dis_out.mean()
-
-        optimizerG.zero_grad()
-        g_loss.backward()
-        optimizerG.step()
-
-        '''
-
-        Train Discriminator: maximize log(D(x)) + log(1 - D(G(z)))
-
-        '''
-
-        for _ in range(num_D_steps):
-
-            if batch_idx+1 == len(train_dataloader):
-                dataloader_iter = iter(train_dataloader)
-                batch_idx = 0
-
-            # get training images
-            batch_train_images, batch_train_labels = dataloader_iter.next()
-            assert batch_size == batch_train_images.shape[0]
-            batch_train_images = batch_train_images.type(torch.float).cuda()
-            batch_train_labels = batch_train_labels.type(torch.long).cuda()
-            batch_idx+=1
-
-            # Measure discriminator's ability to classify real from generated samples
-            if use_DiffAugment:
-                real_dis_out = netD(DiffAugment(batch_train_images, policy=policy), batch_train_labels)
-                fake_dis_out = netD(DiffAugment(batch_fake_images.detach(), policy=policy), batch_train_labels.detach())
-            else:
-                real_dis_out = netD(batch_train_images, batch_train_labels)
-                fake_dis_out = netD(batch_fake_images.detach(), batch_train_labels.detach())
-
-
-            if loss_type == "vanilla":
-                real_dis_out = torch.nn.Sigmoid()(real_dis_out)
-                fake_dis_out = torch.nn.Sigmoid()(fake_dis_out)
-                d_loss_real = - torch.log(real_dis_out+1e-20)
-                d_loss_fake = - torch.log(1-fake_dis_out+1e-20)
-            elif loss_type == "hinge":
-                d_loss_real = torch.nn.ReLU()(1.0 - real_dis_out)
-                d_loss_fake = torch.nn.ReLU()(1.0 + fake_dis_out)
-            d_loss = (d_loss_real + d_loss_fake).mean()
-
-            optimizerD.zero_grad()
-            d_loss.backward()
-            optimizerD.step()
-
-        
-
-        if (niter+1)%20 == 0:
-            print ("cGAN-%s: [Iter %d/%d] [D loss: %.4f] [G loss: %.4f] [D out real:%.4f] [D out fake:%.4f] [Time: %.4f]" % (gan_arch, niter+1, niters, d_loss.item(), g_loss.item(), real_dis_out.mean().item(),fake_dis_out.mean().item(), timeit.default_timer()-start_time))
-
-
-        if (niter+1) % visualize_freq == 0:
-            netG.eval()
-            with torch.no_grad():
-                gen_imgs = netG(z_fixed, y_fixed)
-                gen_imgs = gen_imgs.detach()
-            save_image(gen_imgs.data, save_images_folder +'/{}.png'.format(niter+1), nrow=n_row, normalize=True)
-
-        if save_models_folder is not None and ((niter+1) % save_niters_freq == 0 or (niter+1) == niters):
-            save_file = save_models_folder + "/cGAN_{}_nDsteps_{}_checkpoint_intrain/cGAN_checkpoint_niters_{}.pth".format(gan_arch, num_D_steps, niter+1)
-            os.makedirs(os.path.dirname(save_file), exist_ok=True)
-            torch.save({
-                    'netG_state_dict': netG.state_dict(),
-                    'netD_state_dict': netD.state_dict(),
-                    'optimizerG_state_dict': optimizerG.state_dict(),
-                    'optimizerD_state_dict': optimizerD.state_dict(),
-                    'rng_state': torch.get_rng_state()
-            }, save_file)
-    #end for niter
-
-
-    return netG, netD
-
-
-def sample_cgan_given_labels(netG, given_labels, class_cutoff_points, batch_size = 500, denorm=True, to_numpy=True, verbose=True):
-    '''
-    given_labels: a numpy array; raw label without any normalization; not class label
-    class_cutoff_points: the cutoff points to determine the membership of a give label
-    '''
-
-    class_cutoff_points = np.array(class_cutoff_points)
-    num_classes = len(class_cutoff_points)-1
-
-    nfake = len(given_labels)
-    given_class_labels = np.zeros(nfake)
-    for i in range(nfake):
-        curr_given_label = given_labels[i]
-        diff_tmp = class_cutoff_points - curr_given_label
-        indx_nonneg = np.where(diff_tmp>=0)[0]
-        if len(indx_nonneg)==1: #the last element of diff_tmp is non-negative
-            curr_given_class_label = num_classes-1
-            assert indx_nonneg[0] == num_classes
-        elif len(indx_nonneg)>1:
-            if diff_tmp[indx_nonneg[0]]>0:
-                curr_given_class_label = indx_nonneg[0] - 1
-            else:
-                curr_given_class_label = indx_nonneg[0]
-        given_class_labels[i] = curr_given_class_label
-    given_class_labels = np.concatenate((given_class_labels, given_class_labels[0:batch_size]))
-
-    if batch_size>nfake:
-        batch_size = nfake
-    fake_images = []
-    netG=netG.cuda()
-    netG.eval()
-    with torch.no_grad():
-        if verbose:
-            pb = SimpleProgressBar()
-        tmp = 0
-        while tmp < nfake:
-            z = torch.randn(batch_size, dim_gan, dtype=torch.float).cuda()
-            labels = torch.from_numpy(given_class_labels[tmp:(tmp+batch_size)]).type(torch.long).cuda()
-            if labels.max().item()>num_classes:
-                print("Error: max label {}".format(labels.max().item()))
-            batch_fake_images = netG(z, labels)
-            if denorm: #denorm imgs to save memory
-                assert batch_fake_images.max().item()<=1.0 and batch_fake_images.min().item()>=-1.0
-                batch_fake_images = batch_fake_images*0.5+0.5
-                batch_fake_images = batch_fake_images*255.0
-                batch_fake_images = batch_fake_images.type(torch.uint8)
-                # assert batch_fake_images.max().item()>1
-            fake_images.append(batch_fake_images.detach().cpu())
-            tmp += batch_size
-            if verbose:
-                pb.update(min(float(tmp)/nfake, 1)*100)
-
-    fake_images = torch.cat(fake_images, dim=0)
-    #remove extra entries
-    fake_images = fake_images[0:nfake]
-
-    if to_numpy:
-        fake_images = fake_images.numpy()
-
-    return fake_images, given_labels

myosotis_researches/CcGAN/train_128/train_cgan_concat.py

@@ -1,242 +0,0 @@
-
-import torch
-import torch.nn as nn
-from torchvision.utils import save_image
-import numpy as np
-import os
-import timeit
-
-from .utils import IMGs_dataset, SimpleProgressBar
-from .opts import parse_opts
-from .DiffAugment_pytorch import DiffAugment
-
-''' Settings '''
-args = parse_opts()
-
-# some parameters in opts
-loss_type = args.loss_type_gan
-niters = args.niters_gan
-resume_niters = args.resume_niters_gan
-dim_gan = args.dim_gan
-lr_g = args.lr_g_gan
-lr_d = args.lr_d_gan
-save_niters_freq = args.save_niters_freq
-batch_size = min(args.batch_size_disc, args.batch_size_gene)
-num_classes = args.cGAN_num_classes
-gan_arch = args.GAN_arch
-num_D_steps = args.num_D_steps
-
-visualize_freq = args.visualize_freq
-
-num_workers = args.num_workers
-
-NC = args.num_channels
-IMG_SIZE = args.img_size
-max_label = args.max_label
-
-use_DiffAugment = args.gan_DiffAugment
-policy = args.gan_DiffAugment_policy
-
-
-def train_cgan_concat(images, labels, netG, netD, save_images_folder, save_models_folder = None):
-
-    netG = netG.cuda()
-    netD = netD.cuda()
-
-    optimizerG = torch.optim.Adam(netG.parameters(), lr=lr_g, betas=(0.5, 0.999))
-    optimizerD = torch.optim.Adam(netD.parameters(), lr=lr_d, betas=(0.5, 0.999))
-
-    trainset = IMGs_dataset(images, labels, normalize=True)
-    train_dataloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
-    unique_labels = np.sort(np.array(list(set(labels)))).astype(np.int)
-
-    if save_models_folder is not None and resume_niters>0:
-        save_file = save_models_folder + "/cGAN_{}_nDsteps_{}_checkpoint_intrain/cGAN_checkpoint_niters_{}.pth".format(gan_arch, num_D_steps, resume_niters)
-        checkpoint = torch.load(save_file)
-        netG.load_state_dict(checkpoint['netG_state_dict'])
-        netD.load_state_dict(checkpoint['netD_state_dict'])
-        optimizerG.load_state_dict(checkpoint['optimizerG_state_dict'])
-        optimizerD.load_state_dict(checkpoint['optimizerD_state_dict'])
-        torch.set_rng_state(checkpoint['rng_state'])
-    #end if
-
-    # printed images with labels between the 5-th quantile and 95-th quantile of training labels
-    n_row=10; n_col = n_row
-    z_fixed = torch.randn(n_row*n_col, dim_gan, dtype=torch.float).cuda()
-    start_label = np.quantile(labels, 0.05)
-    end_label = np.quantile(labels, 0.95)
-    selected_labels = np.linspace(start_label, end_label, num=n_row)
-    y_fixed = np.zeros(n_row*n_col)
-    for i in range(n_row):
-        curr_label = selected_labels[i]
-        for j in range(n_col):
-            y_fixed[i*n_col+j] = curr_label
-    print(y_fixed)
-    y_fixed = torch.from_numpy(y_fixed).type(torch.float).view(-1,1).cuda()
-
-
-    batch_idx = 0
-    dataloader_iter = iter(train_dataloader)
-
-    start_time = timeit.default_timer()
-    for niter in range(resume_niters, niters):
-
-        if batch_idx+1 == len(train_dataloader):
-            dataloader_iter = iter(train_dataloader)
-            batch_idx = 0
-
-        '''
-
-        Train Generator: maximize log(D(G(z)))
-
-        '''
-
-        netG.train()
-
-        # get training images
-        _, batch_train_labels = dataloader_iter.next()
-        assert batch_size == batch_train_labels.shape[0]
-        batch_train_labels = batch_train_labels.type(torch.long).cuda()
-        batch_idx+=1
-
-        # Sample noise and labels as generator input
-        z = torch.randn(batch_size, dim_gan, dtype=torch.float).cuda()
-
-        #generate fake images
-        batch_fake_images = netG(z, batch_train_labels)
-
-        # Loss measures generator's ability to fool the discriminator
-        if use_DiffAugment:
-            dis_out = netD(DiffAugment(batch_fake_images, policy=policy), batch_train_labels)
-        else:
-            dis_out = netD(batch_fake_images, batch_train_labels)
-        
-        if loss_type == "vanilla":
-            dis_out = torch.nn.Sigmoid()(dis_out)
-            g_loss = - torch.mean(torch.log(dis_out+1e-20))
-        elif loss_type == "hinge":
-            g_loss = - dis_out.mean()
-
-        optimizerG.zero_grad()
-        g_loss.backward()
-        optimizerG.step()
-
-        '''
-
-        Train Discriminator: maximize log(D(x)) + log(1 - D(G(z)))
-
-        '''
-
-        for _ in range(num_D_steps):
-
-            if batch_idx+1 == len(train_dataloader):
-                dataloader_iter = iter(train_dataloader)
-                batch_idx = 0
-
-            # get training images
-            batch_train_images, batch_train_labels = dataloader_iter.next()
-            assert batch_size == batch_train_images.shape[0]
-            batch_train_images = batch_train_images.type(torch.float).cuda()
-            batch_train_labels = batch_train_labels.type(torch.long).cuda()
-            batch_idx+=1
-
-            # Measure discriminator's ability to classify real from generated samples
-            if use_DiffAugment:
-                real_dis_out = netD(DiffAugment(batch_train_images, policy=policy), batch_train_labels)
-                fake_dis_out = netD(DiffAugment(batch_fake_images.detach(), policy=policy), batch_train_labels.detach())
-            else:
-                real_dis_out = netD(batch_train_images, batch_train_labels)
-                fake_dis_out = netD(batch_fake_images.detach(), batch_train_labels.detach())
-
-
-            if loss_type == "vanilla":
-                real_dis_out = torch.nn.Sigmoid()(real_dis_out)
-                fake_dis_out = torch.nn.Sigmoid()(fake_dis_out)
-                d_loss_real = - torch.log(real_dis_out+1e-20)
-                d_loss_fake = - torch.log(1-fake_dis_out+1e-20)
-            elif loss_type == "hinge":
-                d_loss_real = torch.nn.ReLU()(1.0 - real_dis_out)
-                d_loss_fake = torch.nn.ReLU()(1.0 + fake_dis_out)
-            d_loss = (d_loss_real + d_loss_fake).mean()
-
-            optimizerD.zero_grad()
-            d_loss.backward()
-            optimizerD.step()
-
-        
-
-        if (niter+1)%20 == 0:
-            print ("cGAN(concat)-%s: [Iter %d/%d] [D loss: %.4f] [G loss: %.4f] [D out real:%.4f] [D out fake:%.4f] [Time: %.4f]" % (gan_arch, niter+1, niters, d_loss.item(), g_loss.item(), real_dis_out.mean().item(),fake_dis_out.mean().item(), timeit.default_timer()-start_time))
-
-
-        if (niter+1) % visualize_freq == 0:
-            netG.eval()
-            with torch.no_grad():
-                gen_imgs = netG(z_fixed, y_fixed)
-                gen_imgs = gen_imgs.detach()
-            save_image(gen_imgs.data, save_images_folder +'/{}.png'.format(niter+1), nrow=n_row, normalize=True)
-
-        if save_models_folder is not None and ((niter+1) % save_niters_freq == 0 or (niter+1) == niters):
-            save_file = save_models_folder + "/cGAN_{}_nDsteps_{}_checkpoint_intrain/cGAN_checkpoint_niters_{}.pth".format(gan_arch, num_D_steps, niter+1)
-            os.makedirs(os.path.dirname(save_file), exist_ok=True)
-            torch.save({
-                    'netG_state_dict': netG.state_dict(),
-                    'netD_state_dict': netD.state_dict(),
-                    'optimizerG_state_dict': optimizerG.state_dict(),
-                    'optimizerD_state_dict': optimizerD.state_dict(),
-                    'rng_state': torch.get_rng_state()
-            }, save_file)
-    #end for niter
-
-
-    return netG, netD
-
-
-def sample_cgan_concat_given_labels(netG, given_labels, batch_size = 100, denorm=True, to_numpy=True, verbose=True):
-    '''
-    netG: pretrained generator network
-    given_labels: float. unnormalized labels. we need to convert them to values in [-1,1]. 
-    '''
-
-    ## num of fake images will be generated
-    nfake = len(given_labels)
-
-    ## normalize regression
-    labels = given_labels/max_label
-
-    ## generate images
-    if batch_size>nfake:
-        batch_size = nfake
-
-    fake_images = []
-    ## concat to avoid out of index errors
-    labels = np.concatenate((labels, labels[0:batch_size]), axis=0)
-
-    netG=netG.cuda()
-    netG.eval()
-    with torch.no_grad():
-        if verbose:
-            pb = SimpleProgressBar()
-        tmp = 0
-        while tmp < nfake:
-            z = torch.randn(batch_size, dim_gan, dtype=torch.float).cuda()
-            c = torch.from_numpy(labels[tmp:(tmp+batch_size)]).type(torch.float).cuda()
-            batch_fake_images = netG(z, c)
-            if denorm: #denorm imgs to save memory
-                assert batch_fake_images.max().item()<=1.0 and batch_fake_images.min().item()>=-1.0
-                batch_fake_images = batch_fake_images*0.5+0.5
-                batch_fake_images = batch_fake_images*255.0
-                batch_fake_images = batch_fake_images.type(torch.uint8)
-            fake_images.append(batch_fake_images.detach().cpu())
-            tmp += batch_size
-            if verbose:
-                pb.update(min(float(tmp)/nfake, 1)*100)
-
-    fake_images = torch.cat(fake_images, dim=0)
-    #remove extra entries
-    fake_images = fake_images[0:nfake]
-
-    if to_numpy:
-        fake_images = fake_images.numpy()
-
-    return fake_images, given_labels

myosotis_researches/CcGAN/train_128/utils.py

@@ -1,120 +0,0 @@
-"""
-Some helpful functions
-
-"""
-import numpy as np
-import torch
-import torch.nn as nn
-import torchvision
-import matplotlib.pyplot as plt
-import matplotlib as mpl
-from torch.nn import functional as F
-import sys
-import PIL
-from PIL import Image
-
-
-
-# ################################################################################
-# Progress Bar
-class SimpleProgressBar():
-    def __init__(self, width=50):
-        self.last_x = -1
-        self.width = width
-
-    def update(self, x):
-        assert 0 <= x <= 100 # `x`: progress in percent ( between 0 and 100)
-        if self.last_x == int(x): return
-        self.last_x = int(x)
-        pointer = int(self.width * (x / 100.0))
-        sys.stdout.write( '\r%d%% [%s]' % (int(x), '#' * pointer + '.' * (self.width - pointer)))
-        sys.stdout.flush()
-        if x == 100:
-            print('')
-
-
-
-################################################################################
-# torch dataset from numpy array
-class IMGs_dataset(torch.utils.data.Dataset):
-    def __init__(self, images, labels=None, normalize=False):
-        super(IMGs_dataset, self).__init__()
-
-        self.images = images
-        self.n_images = len(self.images)
-        self.labels = labels
-        if labels is not None:
-            if len(self.images) != len(self.labels):
-                raise Exception('images (' +  str(len(self.images)) +') and labels ('+str(len(self.labels))+') do not have the same length!!!')
-        self.normalize = normalize
-
-
-    def __getitem__(self, index):
-
-        image = self.images[index]
-
-        if self.normalize:
-            image = image/255.0
-            image = (image-0.5)/0.5
-
-        if self.labels is not None:
-            label = self.labels[index]
-            return (image, label)
-        else:
-            return image
-
-    def __len__(self):
-        return self.n_images
-
-
-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

myosotis_researches/CcGAN/train_128_output_10/DiffAugment_pytorch.py

@@ -1,76 +0,0 @@
-# Differentiable Augmentation for Data-Efficient GAN Training
-# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han
-# https://arxiv.org/pdf/2006.10738
-
-import torch
-import torch.nn.functional as F
-
-
-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],
-}