- memory efficient EMA

- fix gradient checkpointing
- evaluate using image reward paper
- use attn_resolution in test

EMA.py

@@ -15,13 +15,14 @@ from torch.optim import Optimizer
 
 
 class EMA(Optimizer):
-    def __init__(self, opt, ema_decay):
+    def __init__(self, opt, ema_decay, memory_efficient=False):
         self.ema_decay = ema_decay
         self.apply_ema = self.ema_decay > 0.
         self.optimizer = opt
         self.state = opt.state
         self.param_groups = opt.param_groups
         self.defaults = {}
+        self.memory_efficient = memory_efficient
 
     def step(self, *args, **kwargs):
         # for group in self.optimizer.param_groups:
@@ -53,11 +54,19 @@ class EMA(Optimizer):
 
                 params[p.shape]['data'].append(p.data)
                 ema[p.shape].append(state['ema'])
+            
+            # def stack(d, dim=0):
+                # return torch.stack([di.cpu() for di in d], dim=dim).cuda()
 
             for i in params:
-                params[i]['data'] = torch.stack(params[i]['data'], dim=0)
-                ema[i] = torch.stack(ema[i], dim=0)
-                ema[i].mul_(self.ema_decay).add_(params[i]['data'], alpha=1. - self.ema_decay)
+                if self.memory_efficient:
+                    for j in range(len(params[i]['data'])):
+                        ema[i][j].mul_(self.ema_decay).add_(params[i]['data'][j], alpha=1. - self.ema_decay)
+                    ema[i] = torch.stack(ema[i], dim=0)
+                else:
+                    params[i]['data'] = torch.stack(params[i]['data'], dim=0)
+                    ema[i] = torch.stack(ema[i], dim=0)
+                    ema[i].mul_(self.ema_decay).add_(params[i]['data'], alpha=1. - self.ema_decay)
 
             for p in group['params']:
                 if p.grad is None:

run.py

@@ -274,10 +274,30 @@ def ddgan_ddb_v7():
     cfg = ddgan_ddb_v1()
     return cfg
 
+def ddgan_ddb_v9():
+    cfg = ddgan_ddb_v3()
+    cfg['model']['attn_resolutions'] = '4 8 16 32'
+    return cfg
+
 def ddgan_laion_aesthetic_v15():
     cfg = ddgan_ddb_v3()
     return cfg
 
+def ddgan_ddb_v10():
+    cfg = ddgan_ddb_v9()
+    return cfg
+
+def ddgan_ddb_v11():
+    cfg = ddgan_ddb_v3()
+    cfg['model']['text_encoder'] = "openclip/ViT-g-14/laion2B-s12B-b42K" 
+    return cfg
+
+def ddgan_ddb_v12():
+    cfg = ddgan_ddb_v3()
+    cfg['model']['text_encoder'] = "openclip/ViT-bigG-14/laion2b_s39b_b160k"
+    return cfg
+
+
 models = [
     ddgan_cifar10_cond17, # cifar10, cross attn for discr
     ddgan_cifar10_cond18, # cifar10, xl encoder
@@ -326,6 +346,10 @@ models = [
     ddgan_ddb_v5,
     ddgan_ddb_v6,
     ddgan_ddb_v7,
+    ddgan_ddb_v9,
+    ddgan_ddb_v10,
+    ddgan_ddb_v11,
+    ddgan_ddb_v12,
 ]
 
 def get_model(model_name):
@@ -334,7 +358,7 @@ def get_model(model_name):
             return model()
 
 
-def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guidance_scale:float=0, fid=False, real_img_dir="", q=0.0, seed=0, nb_images_for_fid=0, scale_factor_h=1, scale_factor_w=1, compute_clip_score=False, eval_name="", scale_method="convolutional"):
+def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guidance_scale:float=0, fid=False, real_img_dir="", q=0.0, seed=0, nb_images_for_fid=0, scale_factor_h=1, scale_factor_w=1, compute_clip_score=False, eval_name="", scale_method="convolutional", compute_image_reward=False):
 
     cfg = get_model(model_name)
     model = cfg['model']
@@ -365,12 +389,16 @@ def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guida
     args['scale_factor_w'] = scale_factor_w
     args['n_mlp'] = model.get("n_mlp")
     args['scale_method'] = scale_method
+    args['attn_resolutions'] = model.get("attn_resolutions", "16")
     if fid:
         args['compute_fid'] = ''
         args['real_img_dir'] = real_img_dir 
         args['nb_images_for_fid'] = nb_images_for_fid
     if compute_clip_score:
         args['compute_clip_score'] = ""
+    
+    if compute_image_reward:
+        args['compute_image_reward'] = ""
     if eval_name:
         args["eval_name"] = eval_name
     cmd = "python -u test_ddgan.py " + " ".join(f"--{k} {v}" for k, v in args.items() if v is not None)

score_sde/models/ncsnpp_generator_adagn.py

@@ -37,6 +37,11 @@ import functools
 import torch
 import numpy as np
 
+try:
+  from fairscale.nn.checkpoint import checkpoint_wrapper
+except Exception:
+  checkpoint_wrapper = lambda x:x
+
 
 ResnetBlockDDPM = layerspp.ResnetBlockDDPMpp_Adagn
 ResnetBlockBigGAN = layerspp.ResnetBlockBigGANpp_Adagn
@@ -63,6 +68,7 @@ class NCSNpp(nn.Module):
   def __init__(self, config):
     super().__init__()
     self.config = config
+    self.grad_checkpointing = config.grad_checkpointing if hasattr(config, "grad_checkpointing") else False
     self.not_use_tanh = config.not_use_tanh
     self.act = act = nn.SiLU()
     self.z_emb_dim = z_emb_dim = config.z_emb_dim
@@ -176,6 +182,8 @@ class NCSNpp(nn.Module):
       raise ValueError(f'resblock type {resblock_type} unrecognized.')
 
     # Downsampling block
+    def wrap(block):
+      return checkpoint_wrapper(block) if self.grad_checkpointing else block
 
     channels = config.num_channels
     if progressive_input != 'none':
@@ -189,18 +197,18 @@ class NCSNpp(nn.Module):
       # Residual blocks for this resolution
       for i_block in range(num_res_blocks):
         out_ch = nf * ch_mult[i_level]
-        modules.append(ResnetBlock(in_ch=in_ch, out_ch=out_ch))
+        modules.append(wrap(ResnetBlock(in_ch=in_ch, out_ch=out_ch)))
         in_ch = out_ch
 
         if all_resolutions[i_level] in attn_resolutions:
-          modules.append(AttnBlock(channels=in_ch))
+          modules.append(wrap(AttnBlock(channels=in_ch)))
         hs_c.append(in_ch)
 
       if i_level != num_resolutions - 1:
         if resblock_type == 'ddpm':
           modules.append(Downsample(in_ch=in_ch))
         else:
-          modules.append(ResnetBlock(down=True, in_ch=in_ch))
+          modules.append(wrap(ResnetBlock(down=True, in_ch=in_ch)))
 
         if progressive_input == 'input_skip':
           modules.append(combiner(dim1=input_pyramid_ch, dim2=in_ch))
@@ -214,21 +222,21 @@ class NCSNpp(nn.Module):
         hs_c.append(in_ch)
 
     in_ch = hs_c[-1]
-    modules.append(ResnetBlock(in_ch=in_ch))
-    modules.append(AttnBlock(channels=in_ch))
-    modules.append(ResnetBlock(in_ch=in_ch))
+    modules.append(wrap(ResnetBlock(in_ch=in_ch)))
+    modules.append(wrap(AttnBlock(channels=in_ch)))
+    modules.append(wrap(ResnetBlock(in_ch=in_ch)))
 
     pyramid_ch = 0
     # Upsampling block
     for i_level in reversed(range(num_resolutions)):
       for i_block in range(num_res_blocks + 1):
         out_ch = nf * ch_mult[i_level]
-        modules.append(ResnetBlock(in_ch=in_ch + hs_c.pop(),
-                                   out_ch=out_ch))
+        modules.append(wrap(ResnetBlock(in_ch=in_ch + hs_c.pop(),
+                                   out_ch=out_ch)))
         in_ch = out_ch
 
       if all_resolutions[i_level] in attn_resolutions:
-        modules.append(AttnBlock(channels=in_ch))
+        modules.append(wrap(AttnBlock(channels=in_ch)))
 
       if progressive != 'none':
         if i_level == num_resolutions - 1:
@@ -260,7 +268,7 @@ class NCSNpp(nn.Module):
         if resblock_type == 'ddpm':
           modules.append(Upsample(in_ch=in_ch))
         else:
-          modules.append(ResnetBlock(in_ch=in_ch, up=True))
+          modules.append(wrap(ResnetBlock(in_ch=in_ch, up=True)))
 
     assert not hs_c
 

test_ddgan.py

@@ -380,7 +380,11 @@ def sample_and_test(args):
         epochs = range(1000)
     else:
         epochs = [args.epoch_id]
-    
+    if args.compute_image_reward:
+        import ImageReward as RM
+        #image_reward = RM.load("ImageReward-v1.0", download_root=".").to(device)
+        image_reward = RM.load("ImageReward.pt", download_root=".").to(device)
+
     for epoch in epochs:
         args.epoch_id = epoch
         path = './saved_info/dd_gan/{}/{}/netG_{}.pth'.format(args.dataset, args.exp, args.epoch_id)
@@ -389,7 +393,7 @@ def sample_and_test(args):
             continue
         if not os.path.exists(next_next_path):
             break
-        print(path)
+        print("PATH", path)
 
         #if not os.path.exists(next_path):
         #    print(f"STOP at {epoch}")
@@ -400,9 +404,7 @@ def sample_and_test(args):
             continue
         suffix = '_' + args.eval_name if args.eval_name else ""
         dest = './saved_info/dd_gan/{}/{}/eval_{}{}.json'.format(args.dataset, args.exp, args.epoch_id, suffix)
-        next_dest = './saved_info/dd_gan/{}/{}/eval_{}{}.json'.format(args.dataset, args.exp, args.epoch_id+1, suffix)
-
-        if (args.compute_fid or args.compute_clip_score) and  os.path.exists(dest):
+        if (args.compute_fid or args.compute_clip_score or args.compute_image_reward) and  os.path.exists(dest):
             continue
         print("Eval Epoch", args.epoch_id)
         #loading weights from ddp in single gpu
@@ -424,7 +426,8 @@ def sample_and_test(args):
         if not os.path.exists(save_dir):
             os.makedirs(save_dir)
         
-        if args.compute_fid or args.compute_clip_score:
+
+        if args.compute_fid or args.compute_clip_score or args.compute_image_reward:
             from torch.nn.functional import adaptive_avg_pool2d
             from pytorch_fid.fid_score import calculate_activation_statistics, calculate_fid_given_paths, ImagePathDataset, compute_statistics_of_path, calculate_frechet_distance
             from pytorch_fid.inception import InceptionV3
@@ -472,6 +475,8 @@ def sample_and_test(args):
             
             if args.compute_clip_score:
                 clip_scores = []
+            if args.compute_image_reward:
+                image_rewards = []
             
             for b in range(0, len(texts), args.batch_size):
                 text = texts[b:b+args.batch_size]
@@ -485,12 +490,7 @@ def sample_and_test(args):
                     else:
                         fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, cond=cond)
                     fake_sample = to_range_0_1(fake_sample)
-                    """
-                    for j, x in enumerate(fake_sample):
-                        index = i * args.batch_size + j 
-                        torchvision.utils.save_image(x, './generated_samples/{}/{}.jpg'.format(args.dataset, index))
-                    """
-
+               
                     if args.compute_fid:
                         with torch.no_grad():
                             pred = inceptionv3(fake_sample)[0]
@@ -511,9 +511,18 @@ def sample_and_test(args):
                             imf = torch.nn.functional.normalize(imf, dim=1)
                             txtf = torch.nn.functional.normalize(txtf, dim=1)
                             clip_scores.append(((imf * txtf).sum(dim=1)).cpu())
-                    
+
+                    if args.compute_image_reward:
+                        for k, sample in enumerate(fake_sample):
+                            img = sample.cpu().numpy().transpose(1,2,0)
+                            img = img * 255
+                            img = img.astype(np.uint8)
+                            text_k = text[k]
+                            score = image_reward.score(text_k, img)
+                            image_rewards.append(score)
                     if i % 10 == 0:
                         print('evaluating batch ', i, time.time() - t0)
+                    #break
                 i += 1
 
             results = {}
@@ -526,6 +535,9 @@ def sample_and_test(args):
             if args.compute_clip_score:
                 clip_score = torch.cat(clip_scores).mean().item()
                 results['clip_score'] = clip_score
+            if args.compute_image_reward:
+                reward = np.mean(image_rewards)
+                results['image_reward'] = reward
             results.update(vars(args))
             with open(dest, "w") as fd:
                 json.dump(results, fd)
@@ -591,6 +603,9 @@ if __name__ == '__main__':
                             help='whether or not compute FID')
     parser.add_argument('--compute_clip_score', action='store_true', default=False,
                             help='whether or not compute CLIP score')
+    parser.add_argument('--compute_image_reward', action='store_true', default=False,
+                            help='whether or not compute CLIP score')
+
     parser.add_argument('--clip_model', type=str,default="ViT-L/14")
     parser.add_argument('--eval_name', type=str,default="")
 
@@ -625,7 +640,7 @@ if __name__ == '__main__':
 
     parser.add_argument('--num_res_blocks', type=int, default=2,
                             help='number of resnet blocks per scale')
-    parser.add_argument('--attn_resolutions', default=(16,),
+    parser.add_argument('--attn_resolutions', default=(16,), nargs='+', type=int,
                             help='resolution of applying attention')
     parser.add_argument('--dropout', type=float, default=0.,
                             help='drop-out rate')

train_ddgan.py

@@ -4,14 +4,14 @@
 # This work is licensed under the NVIDIA Source Code License
 # for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
 # ---------------------------------------------------------------
+import torch
 
 from glob import glob
 import argparse
-import torch
 import numpy as np
-
+import json
 import os
-
+import time
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
@@ -288,6 +288,15 @@ def train(rank, gpu, args):
                     transforms.ToTensor(),
                     transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
             ])
+        elif args.preprocessing == "simple_random_crop_v2":
+            train_transform = transforms.Compose([
+                    transforms.Resize(args.image_size),
+                    transforms.RandomCrop(args.image_size, interpolation=3),
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+            ])
+        else:
+            raise ValueError(args.preprocessing)
         shards = glob(os.path.join(args.dataset_root, "*.tar")) if os.path.isdir(args.dataset_root)  else args.dataset_root
         pipeline = [ResampledShards2(shards)]
         pipeline.extend([
@@ -312,7 +321,7 @@ def train(rank, gpu, args):
             dataset,
             batch_size=None,
             shuffle=False,
-            num_workers=8,
+            num_workers=1,
         )
     
     if args.dataset != "wds":
@@ -355,6 +364,7 @@ def train(rank, gpu, args):
                                 cond_size=text_encoder.output_size,
                                 act=nn.LeakyReLU(0.2)).to(device)
     elif args.discr_type == "large_attn_pool":
+        # Discriminator with  Attention Pool based discriminator for text conditioning
         netD = Discriminator_large(nc = 2*args.num_channels, ngf = args.ngf, 
                                 t_emb_dim = args.t_emb_dim,
                                 cond_size=text_encoder.output_size,
@@ -362,6 +372,7 @@ def train(rank, gpu, args):
                                 act=nn.LeakyReLU(0.2)).to(device)
 
     elif args.discr_type == "large_cond_attn":
+        # Discriminator with  Cross-Attention based discriminator for text conditioning
         netD = CondAttnDiscriminator(
             nc = 2*args.num_channels, 
             ngf = args.ngf, 
@@ -391,7 +402,7 @@ def train(rank, gpu, args):
     optimizerG = optim.Adam(netG.parameters(), lr=args.lr_g, betas = (args.beta1, args.beta2))
     
     if args.use_ema:
-        optimizerG = EMA(optimizerG, ema_decay=args.ema_decay)
+        optimizerG = EMA(optimizerG, ema_decay=args.ema_decay, memory_efficient=args.grad_checkpointing)
     
     schedulerG = torch.optim.lr_scheduler.CosineAnnealingLR(optimizerG, args.num_epoch, eta_min=1e-5)
     schedulerD = torch.optim.lr_scheduler.CosineAnnealingLR(optimizerD, args.num_epoch, eta_min=1e-5)
@@ -403,12 +414,10 @@ def train(rank, gpu, args):
         netD = nn.parallel.DistributedDataParallel(netD, device_ids=[gpu], find_unused_parameters=args.discr_type=="projected_gan")
         #if args.discr_type == "projected_gan":
         #    netD._set_static_graph()
-
     
-    if args.grad_checkpointing:
-        from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
-        netG = checkpoint_wrapper(netG)
-
+    #if args.grad_checkpointing:
+        #from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
+        #netG = checkpoint_wrapper(netG)
     exp = args.exp
     parent_dir = "./saved_info/dd_gan/{}".format(args.dataset)
 
@@ -442,8 +451,9 @@ def train(rank, gpu, args):
         optimizerD.load_state_dict(checkpoint['optimizerD'])
         schedulerD.load_state_dict(checkpoint['schedulerD'])
         global_step = checkpoint['global_step']
-        print("=> loaded checkpoint (epoch {})"
-                  .format(checkpoint['epoch']))
+        if rank == 0:
+            print("=> loaded checkpoint (epoch {})"
+                      .format(checkpoint['epoch']))
     else:
         global_step, epoch, init_epoch = 0, 0, 0
     use_cond_attn_discr = args.discr_type in ("large_cond_attn", "small_cond_attn", "large_attn_pool", "projected_gan")
@@ -454,6 +464,7 @@ def train(rank, gpu, args):
             train_sampler.set_epoch(epoch)
        
         for iteration, (x, y) in enumerate(data_loader):
+            t0 = time.time()
             #print(x.shape)
             if args.dataset != "wds":
                 y = [str(yi) for yi in y.tolist()]
@@ -631,6 +642,8 @@ def train(rank, gpu, args):
                 if rank == 0:
                     print('epoch {} iteration{}, G Loss: {}, D Loss: {}'.format(epoch,iteration, errG.item(), errD.item()))
                     print('Global step:', global_step)
+                    dt = time.time() - t0
+                    print('Time per iteration: ', dt)
             if iteration % 1000 == 0:
                 x_t_1 = torch.randn_like(real_data)
                 with autocast():
@@ -640,7 +653,8 @@ def train(rank, gpu, args):
                 
                 if args.save_content:
                     dist.barrier()
-                    print('Saving content.')
+                    if rank == 0:
+                        print('Saving content.')
                     def to_cpu(d):
                         for k, v in d.items():
                             d[k] = v.cpu()
@@ -677,6 +691,9 @@ def train(rank, gpu, args):
                                     'optimizerD': optimizerD.state_dict(), 'schedulerD': schedulerD.state_dict()}                    
                             torch.save(content, os.path.join(exp_path, 'content.pth'))
                             torch.save(content, os.path.join(exp_path, 'content_backup.pth'))
+                            state_content = {'epoch': epoch + 1, 'global_step': global_step}
+                            with open(os.path.join(exp_path, 'netG_{}.json'.format(epoch)), "w") as fd:
+                                fd.write(json.dumps(state_content))
                             if args.use_ema:
                                 optimizerG.swap_parameters_with_ema(store_params_in_ema=True)                        
                             torch.save(netG.state_dict(), os.path.join(exp_path, 'netG_{}.pth'.format(epoch)))
@@ -685,40 +702,8 @@ def train(rank, gpu, args):
 
             
         if not args.no_lr_decay:
-            
             schedulerG.step()
             schedulerD.step()
-        """
-        if rank == 0:
-            if epoch % 10 == 0:
-                torchvision.utils.save_image(x_pos_sample, os.path.join(exp_path, 'xpos_epoch_{}.png'.format(epoch)), normalize=True)
-            
-            x_t_1 = torch.randn_like(real_data)
-            with autocast():
-                fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1, T, args, cond=(cond_pooled, cond, cond_mask))
-            torchvision.utils.save_image(fake_sample, os.path.join(exp_path, 'sample_discrete_epoch_{}.png'.format(epoch)), normalize=True)
-            
-            if args.save_content:
-                if epoch % args.save_content_every == 0:
-                    print('Saving content.')
-                    content = {'epoch': epoch + 1, 'global_step': global_step, 'args': args,
-                               'netG_dict': netG.state_dict(), 'optimizerG': optimizerG.state_dict(),
-                               'schedulerG': schedulerG.state_dict(), 'netD_dict': netD.state_dict(),
-                               'optimizerD': optimizerD.state_dict(), 'schedulerD': schedulerD.state_dict()}
-                    
-                    torch.save(content, os.path.join(exp_path, 'content.pth'))
-                    torch.save(content, os.path.join(exp_path, 'content_backup.pth'))
-                
-            if epoch % args.save_ckpt_every == 0:
-                if args.use_ema:
-                    optimizerG.swap_parameters_with_ema(store_params_in_ema=True)
-                    
-                torch.save(netG.state_dict(), os.path.join(exp_path, 'netG_{}.pth'.format(epoch)))
-                if args.use_ema:
-                    optimizerG.swap_parameters_with_ema(store_params_in_ema=True)
-        dist.barrier()
-        """
-
 
 def init_processes(rank, size, fn, args):
     """ Initialize the distributed environment. """
@@ -748,12 +733,12 @@ if __name__ == '__main__':
                         help='seed used for initialization')
     
     parser.add_argument('--resume', action='store_true',default=False)
-    parser.add_argument('--masked_mean', action='store_true',default=False)
-    parser.add_argument('--mismatch_loss', action='store_true',default=False)
+    parser.add_argument('--masked_mean', action='store_true',default=False, help="use masked mean to pool from t5-based text encoder")
+    parser.add_argument('--mismatch_loss', action='store_true',default=False, help="use mismatch loss")
     parser.add_argument('--text_encoder', type=str, default="google/t5-v1_1-base")
-    parser.add_argument('--cross_attention', action='store_true',default=False)
-    parser.add_argument('--fsdp', action='store_true',default=False)
-    parser.add_argument('--grad_checkpointing', action='store_true',default=False)
+    parser.add_argument('--cross_attention', action='store_true',default=False, help="use cross attention in generator")
+    parser.add_argument('--fsdp', action='store_true',default=False, help='use FSDP')
+    parser.add_argument('--grad_checkpointing', action='store_true',default=False, help='use grad checkpointing')
 
     parser.add_argument('--image_size', type=int, default=32,
                             help='size of image')
@@ -767,9 +752,8 @@ if __name__ == '__main__':
     parser.add_argument('--beta_max', type=float, default=20.,
                             help='beta_max for diffusion')
     parser.add_argument('--classifier_free_guidance_proba', type=float, default=0.0)
-    
     parser.add_argument('--num_channels_dae', type=int, default=128,
-                            help='number of initial channels in denosing model')
+                            help='number of initial channels in denosing model generator')
     parser.add_argument('--n_mlp', type=int, default=3,
                             help='number of mlp layers for z')
     parser.add_argument('--ch_mult', nargs='+', type=int,
@@ -825,7 +809,7 @@ if __name__ == '__main__':
     parser.add_argument('--beta2', type=float, default=0.9,
                             help='beta2 for adam')
     parser.add_argument('--no_lr_decay',action='store_true', default=False)
-    parser.add_argument('--grad_penalty_cond', action='store_true',default=False)
+    parser.add_argument('--grad_penalty_cond', action='store_true',default=False, help="cond based grad penalty")
 
     parser.add_argument('--use_ema', action='store_true', default=False,
                             help='use EMA or not')