update

EMA.py

@@ -39,7 +39,6 @@ class EMA(Optimizer):
                 # State initialization
                 if 'ema' not in state:
                     state['ema'] = p.data.clone()
-
                 if p.shape not in params:
                     params[p.shape] = {'idx': 0, 'data': []}
                     ema[p.shape] = []

clip_encoder.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import open_clip
+from einops import rearrange
+
+
+def exists(val):
+    return val is not None
+
+class CLIPEncoder(nn.Module):
+
+    def __init__(self, model, pretrained):
+        super().__init__()
+        self.model = model
+        self.pretrained = pretrained
+        self.model, _, _ = open_clip.create_model_and_transforms(model, pretrained=pretrained)
+        self.output_size = self.model.transformer.width
+
+    def forward(self, texts, return_only_pooled=True):
+        device = next(self.parameters()).device
+        toks = open_clip.tokenize(texts).to(device)
+        x = self.model.token_embedding(toks)  # [batch_size, n_ctx, d_model]
+        x = x + self.model.positional_embedding
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.model.transformer(x, attn_mask=self.model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.model.ln_final(x)
+        mask = (toks!=0)
+        pooled = x[torch.arange(x.shape[0]), toks.argmax(dim=-1)] @ self.model.text_projection
+        if return_only_pooled:
+            return pooled
+        else:
+            return pooled, x, mask
+
+
+
+
+class CLIPImageEncoder(nn.Module):
+
+    def __init__(self, model_type="ViT-B/32"):
+        super().__init__()
+        import clip
+        self.model, preprocess = clip.load(model_type, device="cpu", jit=False)
+        CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
+        CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
+        mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1)
+        std = torch.tensor(CLIP_STD).view(1, 3, 1, 1)
+        self.register_buffer("mean", mean)
+        self.register_buffer("std", std)
+        self.output_size = 512
+
+    def forward_image(self, x):
+        x = torch.nn.functional.interpolate(x, mode='bicubic', size=(224, 224))
+        x = (x-self.mean)/self.std
+        return self.model.encode_image(x)
+
+    def forward_text(self, texts):
+        import clip
+        toks = clip.tokenize(texts, truncate=True).to(self.mean.device)
+        return self.model.encode_text(toks)
+
+
+
+

encoder.py

@@ -0,0 +1,9 @@
+import t5
+import clip_encoder
+
+def build_encoder(name, **kwargs):
+    if name.startswith("google"):
+        return t5.T5Encoder(name=name, **kwargs)
+    elif name.startswith("openclip"):
+        _, model, pretrained = name.split("/")
+        return clip_encoder.CLIPEncoder(model, pretrained)

run.py

@@ -132,6 +132,8 @@ def ddgan_laion_aesthetic_v2():
 def ddgan_laion_aesthetic_v3():
     cfg = ddgan_laion_aesthetic_v1()
     cfg['model']['text_encoder'] = "google/t5-v1_1-xl" 
+    cfg['model']['mismatch_loss'] = ''
+    cfg['model']['grad_penalty_cond'] = ''
     return cfg
 
 def ddgan_laion_aesthetic_v4():
@@ -146,6 +148,85 @@ def ddgan_laion_aesthetic_v5():
     cfg['model']['grad_penalty_cond'] = ''
     return cfg
 
+
+
+def ddgan_laion2b_v1():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['mismatch_loss'] = ''
+    cfg['model']['grad_penalty_cond'] = ''
+    cfg['model']['num_channels_dae'] = 224
+    cfg['model']['batch_size'] = 2
+    cfg['model']['discr_type'] = "large_cond_attn"
+    cfg['model']['preprocessing'] = 'random_resized_crop_v1'
+    return cfg
+
+def ddgan_laion_aesthetic_v6():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['no_lr_decay'] = ''
+    return cfg
+
+
+
+def ddgan_laion_aesthetic_v7():
+    cfg = ddgan_laion_aesthetic_v6()
+    cfg['model']['r1_gamma'] = 5
+    return cfg
+
+
+def ddgan_laion_aesthetic_v8():
+    cfg = ddgan_laion_aesthetic_v6()
+    cfg['model']['num_timesteps'] = 8
+    return cfg
+
+def ddgan_laion_aesthetic_v9():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['num_channels_dae'] = 384
+    return cfg
+
+def ddgan_sd_v1():
+    cfg = ddgan_laion_aesthetic_v3()
+    return cfg
+def ddgan_sd_v2():
+    cfg = ddgan_laion_aesthetic_v3()
+    return cfg
+def ddgan_sd_v3():
+    cfg = ddgan_laion_aesthetic_v3()
+    return cfg
+def ddgan_sd_v4():
+    cfg = ddgan_laion_aesthetic_v3()
+    return cfg
+def ddgan_sd_v5():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['num_timesteps'] = 8
+    return cfg
+def ddgan_sd_v6():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['num_channels_dae'] = 192
+    return cfg
+def ddgan_sd_v7():
+    cfg = ddgan_laion_aesthetic_v3()
+    return cfg
+def ddgan_sd_v8():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['image_size'] = 512
+    return cfg
+def ddgan_laion_aesthetic_v12():
+    cfg = ddgan_laion_aesthetic_v3()
+    return cfg
+def ddgan_laion_aesthetic_v13():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['text_encoder'] = "openclip/ViT-H-14/laion2b_s32b_b79k" 
+    return cfg
+
+def ddgan_laion_aesthetic_v14():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['text_encoder'] = "openclip/ViT-H-14/laion2b_s32b_b79k" 
+    return cfg
+def ddgan_sd_v9():
+    cfg = ddgan_laion_aesthetic_v3()
+    cfg['model']['text_encoder'] = "openclip/ViT-H-14/laion2b_s32b_b79k" 
+    return cfg
+
 models = [
     ddgan_cifar10_cond17, # cifar10, cross attn for discr
     ddgan_cifar10_cond18, # cifar10, xl encoder
@@ -166,6 +247,23 @@ models = [
     ddgan_laion_aesthetic_v3, # like ddgan_laion_aesthetic_v1 but trained from scratch with T5-XL (continue from 23aug with mismatch and grad penalty and random_resized_crop_v1)
     ddgan_laion_aesthetic_v4, # like ddgan_laion_aesthetic_v1 but trained from scratch with OpenAI's ClipEncoder 
     ddgan_laion_aesthetic_v5, # fine-tune ddgan_laion_aesthetic_v1 with mismatch and cond grad penalty  losses
+    ddgan_laion_aesthetic_v6, # like v3 but without lr decay
+    ddgan_laion_aesthetic_v7, # like v6 but  with r1 gamma of 5 instead of 1, trying to constrain the discr more.
+    ddgan_laion_aesthetic_v8, # like v6 but with 8 timesteps
+    ddgan_laion_aesthetic_v9,
+    ddgan_laion_aesthetic_v12,
+    ddgan_laion_aesthetic_v13,
+    ddgan_laion_aesthetic_v14,
+    ddgan_laion2b_v1,
+    ddgan_sd_v1,
+    ddgan_sd_v2,
+    ddgan_sd_v3,
+    ddgan_sd_v4,
+    ddgan_sd_v5,
+    ddgan_sd_v6,
+    ddgan_sd_v7,
+    ddgan_sd_v8,
+    ddgan_sd_v9,
 ]
 
 def get_model(model_name):
@@ -174,7 +272,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):
+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"):
 
     cfg = get_model(model_name)
     model = cfg['model']
@@ -204,13 +302,15 @@ def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guida
     args['scale_factor_h'] = scale_factor_h
     args['scale_factor_w'] = scale_factor_w
     args['n_mlp'] = model.get("n_mlp")
+    args['scale_method'] = scale_method
     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 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)
     print(cmd)
     call(cmd, shell=True)
@@ -234,4 +334,4 @@ def eval_results(model_name):
 
 if __name__ == "__main__":
     from clize import run
-    run([test, eval_results])
+    run([test, eval_results])

scripts/init.sh

@@ -32,6 +32,21 @@ if [[ "$machine" == juwelsbooster ]]; then
     ml torchvision/0.12.0
     source /p/project/covidnetx/environments/juwels_booster_2022/bin/activate 
 fi
+if [[ "$machine" == hdfml ]]; then
+    export CUDA_VISIBLE_DEVICES=0,1,2,3
+    ml purge
+    ml use $OTHERSTAGES
+    ml Stages/2022
+    ml GCC/11.2.0
+    ml OpenMPI/4.1.2
+    ml CUDA/11.5
+    ml cuDNN/8.3.1.22-CUDA-11.5
+    ml NCCL/2.12.7-1-CUDA-11.5
+    ml PyTorch/1.11-CUDA-11.5
+    ml Horovod/0.24
+    ml torchvision/0.12.0
+    source envs/hdfml/bin/activate
+fi
 if [[ "$machine" == jusuf ]]; then
     echo not supported
 fi

scripts/run_hdfml.sh

@@ -0,0 +1,25 @@
+#!/bin/bash -x
+#SBATCH --account=cstdl
+#SBATCH --nodes=8
+#SBATCH --ntasks-per-node=4
+#SBATCH --cpus-per-task=8
+#SBATCH --time=06:00:00
+#SBATCH --gres=gpu
+#SBATCH --partition=batch
+ml purge
+ml use $OTHERSTAGES
+ml Stages/2022
+ml GCC/11.2.0
+ml OpenMPI/4.1.2
+ml CUDA/11.5
+ml cuDNN/8.3.1.22-CUDA-11.5
+ml NCCL/2.12.7-1-CUDA-11.5
+ml PyTorch/1.11-CUDA-11.5
+ml Horovod/0.24
+ml torchvision/0.12.0
+source envs/hdfml/bin/activate
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+echo "Job id: $SLURM_JOB_ID"
+export TOKENIZERS_PARALLELISM=false
+export NCCL_ASYNC_ERROR_HANDLING=1
+srun python -u $*

scripts/run_jurecadc_ddp.sh

@@ -13,5 +13,8 @@ source scripts/init.sh
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 echo "Job id: $SLURM_JOB_ID"
 export TOKENIZERS_PARALLELISM=false
-export NCCL_ASYNC_ERROR_HANDLING=1
+#export NCCL_ASYNC_ERROR_HANDLING=1
+export NCCL_IB_TIMEOUT=50
+export UCX_RC_TIMEOUT=4s
+export NCCL_IB_RETRY_CNT=10
 srun python -u $*

test_ddgan.py

@@ -86,7 +86,18 @@ class Posterior_Coefficients():
         self.posterior_mean_coef2 = ((1 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1 - self.alphas_cumprod))
         
         self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min=1e-20))
-        
+
+def predict_q_posterior(coefficients, x_0, x_t, t):
+    mean = (
+        extract(coefficients.posterior_mean_coef1, t, x_t.shape) * x_0
+        + extract(coefficients.posterior_mean_coef2, t, x_t.shape) * x_t
+    )
+    var = extract(coefficients.posterior_variance, t, x_t.shape)
+    log_var_clipped = extract(coefficients.posterior_log_variance_clipped, t, x_t.shape)
+    return mean, var, log_var_clipped
+
+
+
 def sample_posterior(coefficients, x_0,x_t, t):
     
     def q_posterior(x_0, x_t, t):
@@ -150,10 +161,10 @@ def sample_from_model_classifier_free_guidance(coefficients, generator, n_time,
             # eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
             eps = eps_uncond * (1 - guidance_scale) + eps_cond * guidance_scale
             x_0 = (1/torch.sqrt(coefficients.alphas_cumprod[i])) * (x - torch.sqrt(1 - coefficients.alphas_cumprod[i]) * eps)
-           
+            #x_0 = x_0_uncond * (1 - guidance_scale) + x_0_cond * guidance_scale
 
             # Dynamic thresholding
-            q = args.dynamic_thresholding_quantile
+            q = opt.dynamic_thresholding_quantile
             #print("Before", x_0.min(), x_0.max())
             if q:
                 shape = x_0.shape
@@ -180,9 +191,174 @@ def sample_from_model_classifier_free_guidance(coefficients, generator, n_time,
     return x
 
 
+def sample_from_model_classifier_free_guidance_convolutional(coefficients, generator, n_time, x_init, T, opt, text_encoder, cond=None, guidance_scale=0, split_input_params=None):
+    x = x_init
+    null = text_encoder([""] * len(x_init), return_only_pooled=False)
+    #latent_z = torch.randn(x.size(0), opt.nz, device=x.device)
+    ks = split_input_params["ks"]  # eg. (128, 128)
+    stride = split_input_params["stride"]  # eg. (64, 64)
+    uf = split_input_params["vqf"]
+    with torch.no_grad():
+        for i in reversed(range(n_time)):
+            t = torch.full((x.size(0),), i, dtype=torch.int64).to(x.device)
+            t_time = t
+            latent_z = torch.randn(x.size(0), opt.nz, device=x.device)
+            
+            fold, unfold, normalization, weighting = get_fold_unfold(x, ks, stride, split_input_params, uf=uf)
+            x = unfold(x)
+            x = x.view((x.shape[0], -1, ks[0], ks[1], x.shape[-1])) 
+            x_new_list = []
+            for j in range(x.shape[-1]):
+                x_0_uncond = generator(x[:,:,:,:,j], t_time, latent_z, cond=null)            
+                x_0_cond = generator(x[:,:,:,:,j], t_time, latent_z, cond=cond)
+
+                eps_uncond = (x[:,:,:,:,j] - torch.sqrt(coefficients.alphas_cumprod[i]) * x_0_uncond) / torch.sqrt(1 - coefficients.alphas_cumprod[i])
+                eps_cond = (x[:,:,:,:,j] - torch.sqrt(coefficients.alphas_cumprod[i]) * x_0_cond) / torch.sqrt(1 - coefficients.alphas_cumprod[i])
+                
+                eps = eps_uncond * (1 - guidance_scale) + eps_cond * guidance_scale
+                x_0 = (1/torch.sqrt(coefficients.alphas_cumprod[i])) * (x[:,:,:,:,j] - torch.sqrt(1 - coefficients.alphas_cumprod[i]) * eps)
+                q = args.dynamic_thresholding_quantile
+                if q:
+                    shape = x_0.shape
+                    x_0_v = x_0.view(shape[0], -1)
+                    d = torch.quantile(torch.abs(x_0_v), q, dim=1, keepdim=True)
+                    d.clamp_(min=1)
+                    x_0_v = x_0_v.clamp(-d, d) / d
+                    x_0 = x_0_v.view(shape)
+                x_new = sample_posterior(coefficients, x_0, x[:,:,:,:,j], t)
+                x_new_list.append(x_new)
+            
+            o = torch.stack(x_new_list, axis=-1) 
+            #o = o * weighting
+            o = o.view((o.shape[0], -1, o.shape[-1])) 
+            decoded = fold(o)
+            decoded = decoded / normalization
+            x = decoded.detach()
+        
+    return x
+
+def sample_from_model_clip_guidance(coefficients, generator, clip_model, n_time, x_init, T, opt, texts, cond=None, guidance_scale=0):
+    x = x_init
+    text_features = torch.nn.functional.normalize(clip_model.forward_text(texts), dim=1)
+    n_time = 16
+    for i in reversed(range(n_time)):
+        t = torch.full((x.size(0),), i%4, dtype=torch.int64).to(x.device)
+        t_time = t            
+        latent_z = torch.randn(x.size(0), opt.nz, device=x.device)
+        x.requires_grad = True
+        x_0 = generator(x, t_time, latent_z, cond=cond)
+        x_new = sample_posterior(coefficients, x_0, x, t)
+        x_new_n = (x_new + 1) / 2
+        image_features = torch.nn.functional.normalize(clip_model.forward_image(x_new_n), dim=1)
+        loss = (image_features*text_features).sum(dim=1).mean()
+        x_grad, = torch.autograd.grad(loss, x)
+        lr = 3000
+        x = x.detach()
+        print(x.min(),x.max(), lr*x_grad.min(), lr*x_grad.max())
+        x += x_grad * lr
+
+        with torch.no_grad():
+            x_0 = generator(x, t_time, latent_z, cond=cond)
+            x_new = sample_posterior(coefficients, x_0, x, t)
+
+        x = x_new.detach()
+        print(i)
+    return x
+
+def meshgrid(h, w):
+    y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
+    x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
+
+    arr = torch.cat([y, x], dim=-1)
+    return arr
+def delta_border(h, w):
+        """
+        :param h: height
+        :param w: width
+        :return: normalized distance to image border,
+         wtith min distance = 0 at border and max dist = 0.5 at image center
+        """
+        lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
+        arr = meshgrid(h, w) / lower_right_corner
+        dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
+        dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
+        edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
+        return edge_dist
+
+def get_weighting(h, w, Ly, Lx, device, split_input_params):
+    weighting = delta_border(h, w)
+    weighting = torch.clip(weighting, split_input_params["clip_min_weight"],
+                            split_input_params["clip_max_weight"], )
+    weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
+
+    if split_input_params["tie_braker"]:
+        L_weighting = delta_border(Ly, Lx)
+        L_weighting = torch.clip(L_weighting,
+                                    split_input_params["clip_min_tie_weight"],
+                                    split_input_params["clip_max_tie_weight"])
+
+        L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
+        weighting = weighting * L_weighting
+    return weighting
+
+def get_fold_unfold(x, kernel_size, stride, split_input_params, uf=1, df=1):  # todo load once not every time, shorten code
+    """
+    :param x: img of size (bs, c, h, w)
+    :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
+    """
+    bs, nc, h, w = x.shape
+
+    # number of crops in image
+    Ly = (h - kernel_size[0]) // stride[0] + 1
+    Lx = (w - kernel_size[1]) // stride[1] + 1
+
+    if uf == 1 and df == 1:
+        fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+        unfold = torch.nn.Unfold(**fold_params)
+
+        fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
+
+        weighting = get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device, split_input_params).to(x.dtype)
+        normalization = fold(weighting).view(1, 1, h, w)  # normalizes the overlap
+        weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
+
+    elif uf > 1 and df == 1:
+        fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+        unfold = torch.nn.Unfold(**fold_params)
+
+        fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
+                            dilation=1, padding=0,
+                            stride=(stride[0] * uf, stride[1] * uf))
+        fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
+
+        weighting = get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device, split_input_params).to(x.dtype)
+        normalization = fold(weighting).view(1, 1, h * uf, w * uf)  # normalizes the overlap
+        weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
+
+    elif df > 1 and uf == 1:
+        fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+        unfold = torch.nn.Unfold(**fold_params)
+
+        fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
+                            dilation=1, padding=0,
+                            stride=(stride[0] // df, stride[1] // df))
+        fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
+
+        weighting = get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device, split_input_params).to(x.dtype)
+        normalization = fold(weighting).view(1, 1, h // df, w // df)  # normalizes the overlap
+        weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
+
+    else:
+        raise NotImplementedError
+
+    return fold, unfold, normalization, weighting
+
+
+
 #%%
 def sample_and_test(args):
     torch.manual_seed(args.seed)
+
     device = 'cuda:0'
     text_encoder  =build_encoder(name=args.text_encoder, masked_mean=args.masked_mean).to(device)
     args.cond_size = text_encoder.output_size
@@ -197,10 +373,9 @@ def sample_and_test(args):
     
     to_range_0_1 = lambda x: (x + 1.) / 2.
 
-    
+    print(vars(args)) 
     netG = NCSNpp(args).to(device)
-    netG.attn_resolutions = [r * args.scale_factor_w for r in netG.attn_resolutions]
-
+     
     if args.epoch_id == -1:
         epochs = range(1000)
     else:
@@ -209,17 +384,27 @@ def sample_and_test(args):
     for epoch in epochs:
         args.epoch_id = epoch
         path = './saved_info/dd_gan/{}/{}/netG_{}.pth'.format(args.dataset, args.exp, args.epoch_id)
+        next_path = './saved_info/dd_gan/{}/{}/netG_{}.pth'.format(args.dataset, args.exp, args.epoch_id+1)
         if not os.path.exists(path):
             continue
+        print(path)
+
+        #if not os.path.exists(next_path):
+        #    print(f"STOP at {epoch}")
+        #    break
         ckpt = torch.load(path, map_location=device)
-        dest = './saved_info/dd_gan/{}/{}/eval_{}.json'.format(args.dataset, args.exp, args.epoch_id)
+        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 and os.path.exists(dest):
+        if (args.compute_fid or args.compute_clip_score) and  os.path.exists(dest):
             continue
         print("Eval Epoch", args.epoch_id)
         #loading weights from ddp in single gpu
+        #print(ckpt.keys())
         for key in list(ckpt.keys()):
-            ckpt[key[7:]] = ckpt.pop(key)
+            if key.startswith("module"):
+                ckpt[key[7:]] = ckpt.pop(key)
         netG.load_state_dict(ckpt)
         netG.eval()
         
@@ -234,7 +419,7 @@ def sample_and_test(args):
         if not os.path.exists(save_dir):
             os.makedirs(save_dir)
         
-        if args.compute_fid:
+        if args.compute_fid or args.compute_clip_score:
             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
@@ -252,9 +437,11 @@ def sample_and_test(args):
             print("Text size:", len(texts))
             #print("Iters:", iters_needed)
             i = 0
-            dims = 2048
-            block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
-            inceptionv3 = InceptionV3([block_idx]).to(device)
+
+            if args.compute_fid:
+                dims = 2048
+                block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+                inceptionv3 = InceptionV3([block_idx]).to(device)
 
             if args.compute_clip_score:
                 import clip
@@ -264,19 +451,20 @@ def sample_and_test(args):
                 clip_mean = torch.Tensor(CLIP_MEAN).view(1,-1,1,1).to(device)
                 clip_std = torch.Tensor(CLIP_STD).view(1,-1,1,1).to(device)
 
-
-            if not args.real_img_dir.endswith("npz"):
-                real_mu, real_sigma = compute_statistics_of_path(
-                    args.real_img_dir, inceptionv3, args.batch_size, dims, device, 
-                    resize=args.image_size,
-                )
-                np.savez("inception_statistics.npz", mu=real_mu, sigma=real_sigma)
-            else:
-                stats = np.load(args.real_img_dir)
-                real_mu = stats['mu']
-                real_sigma = stats['sigma']
-
-            fake_features = []
+            if args.compute_fid:
+                if not args.real_img_dir.endswith("npz"):
+                    real_mu, real_sigma = compute_statistics_of_path(
+                        args.real_img_dir, inceptionv3, args.batch_size, dims, device, 
+                        resize=args.image_size,
+                    )
+                    np.savez("inception_statistics.npz", mu=real_mu, sigma=real_sigma)
+                else:
+                    stats = np.load(args.real_img_dir)
+                    real_mu = stats['mu']
+                    real_sigma = stats['sigma']
+
+                fake_features = []
+            
             if args.compute_clip_score:
                 clip_scores = []
             
@@ -287,7 +475,6 @@ def sample_and_test(args):
                     bs = len(text)
                     t0 = time.time()
                     x_t_1 = torch.randn(bs, args.num_channels,args.image_size, args.image_size).to(device)
-                    #print(x_t_1.shape)
                     if args.guidance_scale:
                         fake_sample = sample_from_model_classifier_free_guidance(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, text_encoder, cond=cond, guidance_scale=args.guidance_scale)
                     else:
@@ -298,45 +485,39 @@ def sample_and_test(args):
                         index = i * args.batch_size + j 
                         torchvision.utils.save_image(x, './generated_samples/{}/{}.jpg'.format(args.dataset, index))
                     """
-                    with torch.no_grad():
-                        pred = inceptionv3(fake_sample)[0]
-                    # If model output is not scalar, apply global spatial average pooling.
-                    # This happens if you choose a dimensionality not equal 2048.
-                    if pred.size(2) != 1 or pred.size(3) != 1:
-                        pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
-                    pred = pred.squeeze(3).squeeze(2).cpu().numpy()
-                    fake_features.append(pred)
+
+                    if args.compute_fid:
+                        with torch.no_grad():
+                            pred = inceptionv3(fake_sample)[0]
+                        # If model output is not scalar, apply global spatial average pooling.
+                        # This happens if you choose a dimensionality not equal 2048.
+                        if pred.size(2) != 1 or pred.size(3) != 1:
+                            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+                        pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+                        fake_features.append(pred)
 
                     if args.compute_clip_score:
                         with torch.no_grad():
                             clip_ims = torch.nn.functional.interpolate(fake_sample, (224, 224), mode="bicubic")
-                            clip_txt = clip.tokenize(text).to(device)
+                            clip_ims = (clip_ims - clip_mean) / clip_std
+                            clip_txt = clip.tokenize(text, truncate=True).to(device)
                             imf = clip_model.encode_image(clip_ims)
                             txtf = clip_model.encode_text(clip_txt)
                             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())
-                    break
+                    
                     if i % 10 == 0:
-                        print('generating batch ', i, time.time() - t0)
-                    """
-                    if i % 10 == 0:
-                        ff = np.concatenate(fake_features)
-                        fake_mu = np.mean(ff, axis=0)
-                        fake_sigma = np.cov(ff, rowvar=False)
-                        fid =  calculate_frechet_distance(real_mu, real_sigma, fake_mu, fake_sigma)
-                        print("FID", fid)
-                    """
+                        print('evaluating batch ', i, time.time() - t0)
                 i += 1
 
-            fake_features = np.concatenate(fake_features)
-            fake_mu = np.mean(fake_features, axis=0)
-            fake_sigma = np.cov(fake_features, rowvar=False)
-            fid =  calculate_frechet_distance(real_mu, real_sigma, fake_mu, fake_sigma)
-            dest = './saved_info/dd_gan/{}/{}/eval_{}.json'.format(args.dataset, args.exp, args.epoch_id)
-            results = {
-                "fid": fid,
-            }
+            results = {}
+            if args.compute_fid:
+                fake_features = np.concatenate(fake_features)
+                fake_mu = np.mean(fake_features, axis=0)
+                fake_sigma = np.cov(fake_features, rowvar=False)
+                fid =  calculate_frechet_distance(real_mu, real_sigma, fake_mu, fake_sigma)
+                results['fid'] = fid
             if args.compute_clip_score:
                 clip_score = torch.cat(clip_scores).mean().item()
                 results['clip_score'] = clip_score
@@ -344,22 +525,54 @@ def sample_and_test(args):
             with open(dest, "w") as fd:
                 json.dump(results, fd)
             print(results)
-        else:
+        else:            
             if args.cond_text.endswith(".txt"):
                 texts = open(args.cond_text).readlines()
                 texts = [t.strip() for t in texts]
             else:
                 texts = [args.cond_text] * args.batch_size
-            cond = text_encoder(texts, return_only_pooled=False)
-            x_t_1 = torch.randn(len(texts), args.num_channels,args.image_size*args.scale_factor_h, args.image_size*args.scale_factor_w).to(device)
-            t0 = time.time()
-            if args.guidance_scale:
-                fake_sample = sample_from_model_classifier_free_guidance(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, text_encoder, cond=cond, guidance_scale=args.guidance_scale)
+            clip_guidance = False
+            if clip_guidance:
+                from clip_encoder import CLIPImageEncoder
+                cond = text_encoder(texts, return_only_pooled=False)
+                clip_image_model = CLIPImageEncoder().to(device)
+                x_t_1 = torch.randn(len(texts), args.num_channels,args.image_size*args.scale_factor_h, args.image_size*args.scale_factor_w).to(device)
+                fake_sample = sample_from_model_clip_guidance(pos_coeff, netG, clip_image_model, args.num_timesteps, x_t_1,T,  args, texts, cond=cond, guidance_scale=args.guidance_scale)
+                fake_sample = to_range_0_1(fake_sample)
+                torchvision.utils.save_image(fake_sample, './samples_{}.jpg'.format(args.dataset))
+
             else:
-                fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, cond=cond)
-            print(time.time() - t0)
-            fake_sample = to_range_0_1(fake_sample)
-            torchvision.utils.save_image(fake_sample, './samples_{}.jpg'.format(args.dataset))
+                cond = text_encoder(texts, return_only_pooled=False)
+                x_t_1 = torch.randn(len(texts), args.num_channels,args.image_size*args.scale_factor_h, args.image_size*args.scale_factor_w).to(device)
+                t0 = time.time()
+                if args.guidance_scale:
+                    if args.scale_factor_h > 1 or args.scale_factor_w > 1:
+                        if args.scale_method == "convolutional":
+                            split_input_params = {
+                                "ks": (args.image_size, args.image_size),
+                                "stride": (150,  150),
+                                "clip_max_tie_weight": 0.5,
+                                "clip_min_tie_weight": 0.01,
+                                "clip_max_weight": 0.5,
+                                "clip_min_weight": 0.01,
+
+                                "tie_braker": True,
+                                'vqf': 1,
+                            }
+                            fake_sample = sample_from_model_classifier_free_guidance_convolutional(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, text_encoder, cond=cond, guidance_scale=args.guidance_scale, split_input_params=split_input_params)
+                        elif args.scale_method == "larger_input":
+                            netG.attn_resolutions = [r * args.scale_factor_w for r in netG.attn_resolutions]
+                            fake_sample = sample_from_model_classifier_free_guidance(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, text_encoder, cond=cond, guidance_scale=args.guidance_scale)
+                    else:
+                        fake_sample = sample_from_model_classifier_free_guidance(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, text_encoder, cond=cond, guidance_scale=args.guidance_scale)
+                else:
+                    fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args, cond=cond)
+
+                print(time.time() - t0)
+                fake_sample = to_range_0_1(fake_sample)
+                torchvision.utils.save_image(fake_sample, './samples_{}.jpg'.format(args.dataset))
+
+
 
     
     
@@ -374,6 +587,7 @@ if __name__ == '__main__':
     parser.add_argument('--compute_clip_score', 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="")
 
     parser.add_argument('--epoch_id', type=int,default=1000)
     parser.add_argument('--guidance_scale', type=float,default=0)
@@ -381,6 +595,8 @@ if __name__ == '__main__':
     parser.add_argument('--cond_text', type=str,default="0")
     parser.add_argument('--scale_factor_h', type=int,default=1)
     parser.add_argument('--scale_factor_w', type=int,default=1)
+    parser.add_argument('--scale_method', type=str,default="convolutional")
+
     parser.add_argument('--cross_attention', action='store_true',default=False)
 
     

train_ddgan.py

@@ -5,7 +5,7 @@
 # for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
 # ---------------------------------------------------------------
 
-
+from glob import glob
 import argparse
 import torch
 import numpy as np
@@ -30,6 +30,7 @@ import shutil
 import logging
 from encoder import build_encoder
 from utils import ResampledShards2
+from torch.utils.tensorboard import SummaryWriter
 
 
 def log_and_continue(exn):
@@ -194,23 +195,29 @@ def sample_from_model(coefficients, generator, n_time, x_init, T, opt, cond=None
         
     return x
 
-
+from contextlib import suppress
 
 def filter_no_caption(sample):
     return 'txt' in sample
 
-
+def get_autocast(precision):
+    if precision == 'amp':
+        return torch.cuda.amp.autocast
+    elif precision == 'amp_bfloat16':
+        return lambda: torch.cuda.amp.autocast(dtype=torch.bfloat16)
+    else:
+        return suppress 
 
 def train(rank, gpu, args):
     from score_sde.models.discriminator import Discriminator_small, Discriminator_large, CondAttnDiscriminator, SmallCondAttnDiscriminator
     from score_sde.models.ncsnpp_generator_adagn import NCSNpp
     from EMA import EMA
     
-    torch.manual_seed(args.seed + rank)
-    torch.cuda.manual_seed(args.seed + rank)
-    torch.cuda.manual_seed_all(args.seed + rank)
+    #torch.manual_seed(args.seed + rank)
+    #torch.cuda.manual_seed(args.seed + rank)
+    #torch.cuda.manual_seed_all(args.seed + rank)
     device = "cuda"
-    
+    autocast = get_autocast(args.precision)
     batch_size = args.batch_size
     
     nz = args.nz #latent dimension
@@ -270,11 +277,12 @@ def train(rank, gpu, args):
             ])
         elif args.preprocessing == "random_resized_crop_v1":
             train_transform = transforms.Compose([
-                    transforms.RandomResizedCrop(256, scale=(0.95, 1.0), interpolation=3),
+                    transforms.RandomResizedCrop(args.image_size, scale=(0.95, 1.0), interpolation=3),
                     transforms.ToTensor(),
                     transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
             ])
-        pipeline = [ResampledShards2(args.dataset_root)]
+        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([
             wds.split_by_node,
             wds.split_by_worker,
@@ -339,6 +347,13 @@ def train(rank, gpu, args):
                                 t_emb_dim = args.t_emb_dim,
                                 cond_size=text_encoder.output_size,
                                 act=nn.LeakyReLU(0.2)).to(device)
+    elif args.discr_type == "large_attn_pool":
+        netD = Discriminator_large(nc = 2*args.num_channels, ngf = args.ngf, 
+                                t_emb_dim = args.t_emb_dim,
+                                cond_size=text_encoder.output_size,
+                                attn_pool=True,
+                                act=nn.LeakyReLU(0.2)).to(device)
+
     elif args.discr_type == "large_cond_attn":
         netD = CondAttnDiscriminator(
             nc = 2*args.num_channels, 
@@ -350,6 +365,15 @@ def train(rank, gpu, args):
     broadcast_params(netG.parameters())
     broadcast_params(netD.parameters())
     
+    if args.fsdp:
+        from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
+        from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+        netG = FSDP(
+            netG,
+            flatten_parameters=True,
+            verbose=True,
+        )
+
     optimizerD = optim.Adam(netD.parameters(), lr=args.lr_d, betas = (args.beta1, args.beta2))
     optimizerG = optim.Adam(netG.parameters(), lr=args.lr_g, betas = (args.beta1, args.beta2))
     
@@ -358,9 +382,16 @@ def train(rank, gpu, args):
     
     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)
+
+    if args.fsdp:   
+        netD = nn.parallel.DistributedDataParallel(netD, device_ids=[gpu])
+    else:
+        netG = nn.parallel.DistributedDataParallel(netG, device_ids=[gpu])
+        netD = nn.parallel.DistributedDataParallel(netD, device_ids=[gpu])
     
-    netG = nn.parallel.DistributedDataParallel(netG, device_ids=[gpu])
-    netD = nn.parallel.DistributedDataParallel(netD, device_ids=[gpu])
+    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)
@@ -377,6 +408,10 @@ def train(rank, gpu, args):
     T = get_time_schedule(args, device)
     
     checkpoint_file = os.path.join(exp_path, 'content.pth')
+    
+    if rank == 0:
+        log_writer = SummaryWriter(exp_path)
+
     if args.resume and os.path.exists(checkpoint_file):
         checkpoint = torch.load(checkpoint_file, map_location="cpu")
         init_epoch = checkpoint['epoch']
@@ -395,7 +430,7 @@ def train(rank, gpu, args):
                   .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")
+    use_cond_attn_discr = args.discr_type in ("large_cond_attn", "small_cond_attn", "large_attn_pool")
     for epoch in range(init_epoch, args.num_epoch+1):
         if args.dataset == "wds":
             os.environ["WDS_EPOCH"] = str(epoch)
@@ -403,6 +438,7 @@ def train(rank, gpu, args):
             train_sampler.set_epoch(epoch)
        
         for iteration, (x, y) in enumerate(data_loader):
+            #print(x.shape)
             if args.dataset != "wds":
                 y = [str(yi) for yi in y.tolist()]
             
@@ -437,15 +473,15 @@ def train(rank, gpu, args):
                     cond_for_discr.requires_grad = True
 
             # train with real
-            D_real = netD(x_t, t, x_tp1.detach(), cond=cond_for_discr).view(-1)
-            
-            errD_real = F.softplus(-D_real)
-            errD_real = errD_real.mean()
+            with autocast():
+                D_real = netD(x_t, t, x_tp1.detach(), cond=cond_for_discr).view(-1)
+                errD_real = F.softplus(-D_real)
+                errD_real = errD_real.mean()
 
             
             errD_real.backward(retain_graph=True)
             
-            
+            grad_penalty = None
             if args.lazy_reg is None:
                 if args.grad_penalty_cond:
                     inputs = (x_t,) + (cond,) if use_cond_attn_discr else (cond_for_discr,)
@@ -491,26 +527,36 @@ def train(rank, gpu, args):
 
             # train with fake
             latent_z = torch.randn(batch_size, nz, device=device)
-            
-            x_0_predict = netG(x_tp1.detach(), t, latent_z, cond=(cond_pooled, cond, cond_mask))
-            x_pos_sample = sample_posterior(pos_coeff, x_0_predict, x_tp1, t)
-            
-            output = netD(x_pos_sample, t, x_tp1.detach(), cond=cond_for_discr).view(-1)
+            with autocast():
+                if args.grad_checkpointing:
+                    ginp  = x_tp1.detach()
+                    ginp.requires_grad = True
+                    latent_z.requires_grad = True
+                    cond_pooled.requires_grad = True
+                    cond.requires_grad = True
+                    #cond_mask.requires_grad = True
+                    x_0_predict = netG(ginp, t, latent_z, cond=(cond_pooled, cond, cond_mask))
+                else:
+                    x_0_predict = netG(x_tp1.detach(), t, latent_z, cond=(cond_pooled, cond, cond_mask))
+                x_pos_sample = sample_posterior(pos_coeff, x_0_predict, x_tp1, t)
                 
-            
-            errD_fake = F.softplus(output)
-            errD_fake = errD_fake.mean()
+                output = netD(x_pos_sample, t, x_tp1.detach(), cond=cond_for_discr).view(-1)
+                    
+                
+                errD_fake = F.softplus(output)
+                errD_fake = errD_fake.mean()
 
             if args.mismatch_loss:
                 # following https://github.com/tobran/DF-GAN/blob/bc38a4f795c294b09b4ef5579cd4ff78807e5b96/code/lib/modules.py,
                 # we add a discr loss for (real image, non matching text)
                 #inds = torch.flip(torch.arange(len(x_t)), dims=(0,))
-                inds = torch.cat([torch.arange(1,len(x_t)),torch.arange(1)])
-                cond_for_discr_mis =  (cond_pooled[inds], cond[inds], cond_mask[inds]) if use_cond_attn_discr else cond_pooled[inds]
-                D_real_mis = netD(x_t, t, x_tp1.detach(), cond=cond_for_discr_mis).view(-1)
-                errD_real_mis = F.softplus(D_real_mis)
-                errD_real_mis = errD_real_mis.mean()
-                errD_fake = errD_fake * 0.5 + errD_real_mis * 0.5
+                with autocast():
+                    inds = torch.cat([torch.arange(1,len(x_t)),torch.arange(1)])
+                    cond_for_discr_mis =  (cond_pooled[inds], cond[inds], cond_mask[inds]) if use_cond_attn_discr else cond_pooled[inds]
+                    D_real_mis = netD(x_t, t, x_tp1.detach(), cond=cond_for_discr_mis).view(-1)
+                    errD_real_mis = F.softplus(D_real_mis)
+                    errD_real_mis = errD_real_mis.mean()
+                    errD_fake = errD_fake * 0.5 + errD_real_mis * 0.5
         
             errD_fake.backward()
     
@@ -534,58 +580,106 @@ def train(rank, gpu, args):
             
             latent_z = torch.randn(batch_size, nz,device=device)
             
-            
+            with autocast():
+                if args.grad_checkpointing:
+                    ginp  = x_tp1.detach()
+                    ginp.requires_grad = True
+                    latent_z.requires_grad = True
+                    cond_pooled.requires_grad = True
+                    cond.requires_grad = True
+                    #cond_mask.requires_grad = True
+                    x_0_predict = netG(ginp, t, latent_z, cond=(cond_pooled, cond, cond_mask))
+                else:
+                    x_0_predict = netG(x_tp1.detach(), t, latent_z, cond=(cond_pooled, cond, cond_mask))
+                x_pos_sample = sample_posterior(pos_coeff, x_0_predict, x_tp1, t)
                 
-            x_0_predict = netG(x_tp1.detach(), t, latent_z, cond=(cond_pooled, cond, cond_mask))
-            x_pos_sample = sample_posterior(pos_coeff, x_0_predict, x_tp1, t)
-            
-            output = netD(x_pos_sample, t, x_tp1.detach(), cond=cond_for_discr).view(-1)
-               
-            
-            errG = F.softplus(-output)
-            errG = errG.mean()
+                output = netD(x_pos_sample, t, x_tp1.detach(), cond=cond_for_discr).view(-1)
+                
+                
+                errG = F.softplus(-output)
+                errG = errG.mean()
             
             errG.backward()
             optimizerG.step()
                 
-           
+            if (iteration % 10 == 0) and (rank == 0):
+                log_writer.add_scalar('g_loss', errG.item(), global_step)
+                log_writer.add_scalar('d_loss', errD.item(), global_step)
+                if grad_penalty is not None:
+                    log_writer.add_scalar('grad_penalty', grad_penalty.item(), global_step)
             
             global_step += 1
+
+
             if iteration % 100 == 0:
                 if rank == 0:
                     print('epoch {} iteration{}, G Loss: {}, D Loss: {}'.format(epoch,iteration, errG.item(), errD.item()))
+                    print('Global step:', global_step)
             if iteration % 1000 == 0:
                 x_t_1 = torch.randn_like(real_data)
-                fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1, T, args, cond=(cond_pooled, cond, cond_mask))
+                with autocast():
+                    fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1, T, args, cond=(cond_pooled, cond, cond_mask))
                 if rank == 0:
                     torchvision.utils.save_image(fake_sample, os.path.join(exp_path, 'sample_discrete_epoch_{}_iteration_{}.png'.format(epoch, iteration)), normalize=True)
-                    if args.save_content:
-                        print('Saving content.')
+                
+                if args.save_content:
+                    dist.barrier()
+                    print('Saving content.')
+                    def to_cpu(d):
+                        for k, v in d.items():
+                            d[k] = v.cpu()
+                        return d
+                    
+                    if args.fsdp:
+                        netG_state_dict = to_cpu(netG.state_dict())
+                        netD_state_dict = to_cpu(netD.state_dict())
+                        #netG_optim_state_dict = (netG.gather_full_optim_state_dict(optimizerG))
+                        netG_optim_state_dict = optimizerG.state_dict()
+                        #print(netG_optim_state_dict)
+                        netD_optim_state_dict = (optimizerD.state_dict())
                         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 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)
+                                'netG_dict': netG_state_dict, 'optimizerG': netG_optim_state_dict,
+                                'schedulerG': schedulerG.state_dict(), 'netD_dict': netD_state_dict,
+                                'optimizerD': netD_optim_state_dict, 'schedulerD': schedulerD.state_dict()}
+                        if rank == 0:
+                            torch.save(content, os.path.join(exp_path, 'content.pth'))
+                            torch.save(content, os.path.join(exp_path, 'content_backup.pth'))
+                        if args.use_ema:
+                            optimizerG.swap_parameters_with_ema(store_params_in_ema=True)                        
+                        if args.use_ema and rank == 0:
+                            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)
+                        #if args.use_ema:
+                        #    dist.barrier()
+                        print("Saved content")
+                    else:
+                        if rank == 0:
+                            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 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)
+
             
         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)
-            fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1, T, args, cond=(cond_pooled, cond, cond_mask))
+            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:
@@ -606,7 +700,8 @@ def train(rank, gpu, args):
                 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):
@@ -641,6 +736,8 @@ if __name__ == '__main__':
     parser.add_argument('--mismatch_loss', action='store_true',default=False)
     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('--image_size', type=int, default=32,
                             help='size of image')
@@ -728,6 +825,7 @@ if __name__ == '__main__':
     parser.add_argument('--save_ckpt_every', type=int, default=25, help='save ckpt every x epochs')
     parser.add_argument('--discr_type', type=str, default="large")
     parser.add_argument('--preprocessing', type=str, default="resize")
+    parser.add_argument('--precision', type=str, default="fp32")
 
     ###ddp
     parser.add_argument('--num_proc_node', type=int, default=1,
@@ -746,4 +844,4 @@ if __name__ == '__main__':
     args.world_size =  int(os.getenv("SLURM_NTASKS"))
     args.rank = int(os.environ['SLURM_PROCID'])
     # size = args.num_process_per_node
-    init_processes(args.rank, args.world_size, train, args)
+    init_processes(args.rank, args.world_size, train, args)

utils.py

@@ -41,7 +41,8 @@ class ResampledShards2(IterableDataset):
         """
         super().__init__()
         #urls = wds.shardlists.expand_urls(urls)
-        urls = list(braceexpand.braceexpand(urls))
+        if type(urls) != list:
+            urls = list(braceexpand.braceexpand(urls))
         self.urls = urls
         assert isinstance(self.urls[0], str)
         self.nshards = nshards