text to image support

pytorch_fid/fid_score.py

@@ -140,7 +140,7 @@ def get_activations(files, model, batch_size=50, dims=2048, device='cpu', resize
                                              batch_size=batch_size,
                                              shuffle=False,
                                              drop_last=False,
-                                             num_workers=cpu_count())
+                                             num_workers=8)
 
     pred_arr = np.empty((len(files), dims))
 
@@ -148,7 +148,7 @@ def get_activations(files, model, batch_size=50, dims=2048, device='cpu', resize
 
     for batch in tqdm(dataloader):
         batch = batch.to(device)
-
+        print(batch.shape, batch.min(), batch.max)
         with torch.no_grad():
             pred = model(batch)[0]
 

run.py

@@ -0,0 +1,127 @@
+import os
+from clize import run
+from glob import glob
+from subprocess import call
+
+def base():
+    return {
+        "slurm":{
+            "t": 360,
+            "N": 2,
+            "n": 8,
+        },
+        "model":{
+            "dataset" :"wds",
+            "dataset_root": "/p/scratch/ccstdl/cherti1/CC12M/{00000..01099}.tar",
+            "image_size": 256,
+            "num_channels": 3,
+            "num_channels_dae": 128,
+            "ch_mult": "1 1 2 2 4 4",
+            "num_timesteps": 4,
+            "num_res_blocks": 2,
+            "batch_size": 8,
+            "num_epoch": 1000,
+            "ngf": 64,
+            "embedding_type": "positional",
+            "use_ema": "",
+            "ema_decay": 0.999,
+            "r1_gamma": 1.0,
+            "z_emb_dim": 256,
+            "lr_d": 1e-4,
+            "lr_g": 1.6e-4,
+            "lazy_reg": 10,
+            "save_content": "",
+            "save_ckpt_every": 1,
+            "masked_mean": "",
+            "resume": "",
+        }
+    }
+def ddgan_cc12m_v2():
+    cfg =  base()
+    cfg['slurm']['N'] = 2
+    cfg['slurm']['n'] = 8
+    return cfg
+
+def ddgan_cc12m_v6():
+    cfg = base()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-large"
+    return cfg
+
+def ddgan_cc12m_v7():
+    cfg = base()
+    cfg['model']['classifier_free_guidance_proba'] = 0.2
+    cfg['slurm']['N'] = 2
+    cfg['slurm']['n'] = 8
+    return cfg
+
+def ddgan_cc12m_v8():
+    cfg = base()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-large"    
+    cfg['model']['classifier_free_guidance_proba'] = 0.2
+    return cfg
+
+def ddgan_cc12m_v9():
+    cfg = base()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-large"    
+    cfg['model']['classifier_free_guidance_proba'] = 0.2
+    cfg['model']['num_channels_dae'] = 320
+    cfg['model']['image_size'] = 64
+    cfg['model']['batch_size'] = 1
+    return cfg
+
+
+def ddgan_cc12m_v11():
+    cfg = base()
+    cfg['model']['text_encoder'] = "google/t5-v1_1-large"    
+    cfg['model']['classifier_free_guidance_proba'] = 0.2
+    cfg['model']['cross_attention'] = ""
+    return cfg
+
+models = [
+    ddgan_cc12m_v2,
+    ddgan_cc12m_v6,
+    ddgan_cc12m_v7,
+    ddgan_cc12m_v8,
+    ddgan_cc12m_v9,
+    ddgan_cc12m_v11,
+
+]
+def get_model(model_name):
+    for model in models:
+        if model.__name__ == 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=""):
+
+    cfg = get_model(model_name)
+    model = cfg['model']
+    if epoch is None:
+        paths = glob('./saved_info/dd_gan/{}/{}/netG_*.pth'.format(model["dataset"], model_name))
+        epoch = max(
+            [int(os.path.basename(path).replace(".pth", "").split("_")[1]) for path in paths]
+        )
+    args = {}
+    args['exp'] = model_name
+    args['image_size'] = model['image_size']
+    args['num_channels'] = model['num_channels']
+    args['dataset'] = model['dataset']
+    args['num_channels_dae'] = model['num_channels_dae']
+    args['ch_mult'] = model['ch_mult']
+    args['num_timesteps'] = model['num_timesteps']
+    args['num_res_blocks'] = model['num_res_blocks']
+    args['batch_size'] = model['batch_size'] if batch_size is None else batch_size
+    args['epoch'] = epoch
+    args['cond_text'] = f'"{cond_text}"'
+    args['text_encoder'] = model.get("text_encoder")
+    args['cross_attention'] = model.get("cross_attention")
+    args['guidance_scale'] = guidance_scale
+
+    if fid:
+        args['compute_fid'] = ''
+        args['real_img_dir'] = real_img_dir 
+    cmd = "python test_ddgan.py " + " ".join(f"--{k} {v}" for k, v in args.items() if v is not None)
+    print(cmd)
+    call(cmd, shell=True)
+
+run([test])

score_sde/models/discriminator.py

@@ -96,11 +96,12 @@ class DownConvBlock(nn.Module):
 class Discriminator_small(nn.Module):
   """A time-dependent discriminator for small images (CIFAR10, StackMNIST)."""
 
-  def __init__(self, nc = 3, ngf = 64, t_emb_dim = 128, act=nn.LeakyReLU(0.2)):
+  def __init__(self, nc = 3, ngf = 64, t_emb_dim = 128, act=nn.LeakyReLU(0.2), cond_size=768):
     super().__init__()
     # Gaussian random feature embedding layer for time
     self.act = act
-    
+    self.cond_proj = nn.Linear(cond_size, ngf*8) 
+    # self.cond_proj.weight.data = default_initializer()(self.cond_proj.weight.shape)
     
     self.t_embed = TimestepEmbedding(
         embedding_dim=t_emb_dim,
@@ -131,10 +132,11 @@ class Discriminator_small(nn.Module):
     self.stddev_feat = 1
     
         
-  def forward(self, x, t, x_t):
-    t_embed = self.act(self.t_embed(t))  
-    
-  
+  def forward(self, x, t, x_t, cond=None):
+    t_embed = self.t_embed(t)
+    # if cond is not None:
+        # t_embed = t_embed + self.cond_proj(cond)
+    t_embed = self.act(t_embed)  
     input_x = torch.cat((x, x_t), dim = 1)
     
     h0 = self.start_conv(input_x)
@@ -159,10 +161,9 @@ class Discriminator_small(nn.Module):
     
     out = self.final_conv(out)
     out = self.act(out)
-   
     
     out = out.view(out.shape[0], out.shape[1], -1).sum(2)
-    out = self.end_linear(out)
+    out = self.end_linear(out) + (self.cond_proj(cond) * out).sum(dim=1, keepdim=True)
     
     return out
 
@@ -170,9 +171,10 @@ class Discriminator_small(nn.Module):
 class Discriminator_large(nn.Module):
   """A time-dependent discriminator for large images (CelebA, LSUN)."""
 
-  def __init__(self, nc = 1, ngf = 32, t_emb_dim = 128, act=nn.LeakyReLU(0.2)):
+  def __init__(self, nc = 1, ngf = 32, t_emb_dim = 128, act=nn.LeakyReLU(0.2), cond_size=768):
     super().__init__()
     # Gaussian random feature embedding layer for time
+    self.cond_proj = nn.Linear(cond_size, ngf*8) 
     self.act = act
     
     self.t_embed = TimestepEmbedding(
@@ -202,8 +204,9 @@ class Discriminator_large(nn.Module):
     self.stddev_feat = 1
     
         
-  def forward(self, x, t, x_t):
-    t_embed = self.act(self.t_embed(t))  
+  def forward(self, x, t, x_t, cond=None):
+    t_embed = self.t_embed(t)
+    t_embed = self.act(t_embed)  
     
     input_x = torch.cat((x, x_t), dim = 1)
     
@@ -233,7 +236,6 @@ class Discriminator_large(nn.Module):
     out = self.act(out)
     
     out = out.view(out.shape[0], out.shape[1], -1).sum(2)
-    out = self.end_linear(out)
-    
+    out = self.end_linear(out) + (self.cond_proj(cond) * out).sum(dim=1, keepdim=True)
     return out
 

score_sde/models/layers.py

@@ -538,6 +538,276 @@ class AttnBlock(nn.Module):
     return x + h
 
 
+class CondAttnBlock(nn.Module):
+  """Channel-wise self-attention block."""
+  def __init__(self, channels, context_dim, dim_head=64, heads=8, norm_context=False, cosine_sim_attn=False):
+    super().__init__()
+    self.GroupNorm_0 = nn.GroupNorm(num_groups=32, num_channels=channels, eps=1e-6)
+    self.ca = CrossAttention(
+      channels,
+      context_dim=context_dim,
+      dim_head=dim_head,
+      heads=heads,
+      norm_context=norm_context,
+      cosine_sim_attn=cosine_sim_attn,
+    )
+  
+  def forward(self, x, cond, mask=None):
+    B, C, H, W = x.shape
+    h = self.GroupNorm_0(x)
+    h = h.view(B, C, H*W)
+    h = h.permute(0,2,1)
+    h = h.contiguous()
+    h_new = self.ca(h, cond, mask=mask)
+    h_new = h_new.permute(0,2,1)
+    h_new = h_new.contiguous()
+    h_new = h_new.view(B, C, H, W)
+    return x + h_new
+  
+from torch import einsum
+from einops import rearrange, repeat, reduce
+from einops.layers.torch import Rearrange, Reduce
+from einops_exts import rearrange_many, repeat_many, check_shape
+from einops_exts.torch import EinopsToAndFrom
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+class Identity(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        context_dim = None,
+        dim_head = 64,
+        heads = 8,
+        norm_context = False,
+        cosine_sim_attn = False
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5 if not cosine_sim_attn else 1.
+        self.cosine_sim_attn = cosine_sim_attn
+        self.cosine_sim_scale = 16 if cosine_sim_attn else 1
+
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        context_dim = default(context_dim, dim)
+
+        self.norm = nn.LayerNorm(dim)
+        self.norm_context = nn.LayerNorm(context_dim) if norm_context else Identity()
+
+        self.null_kv = nn.Parameter(torch.randn(2, dim_head))
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim, bias = False),
+            nn.LayerNorm(dim)
+        )
+
+    def forward(self, x, context, mask = None):
+        b, n, device = *x.shape[:2], x.device
+
+        x = self.norm(x)
+        context = self.norm_context(context)
+
+        q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
+
+        q, k, v = rearrange_many((q, k, v), 'b n (h d) -> b h n d', h = self.heads)
+
+        # add null key / value for classifier free guidance in prior net
+
+        nk, nv = repeat_many(self.null_kv.unbind(dim = -2), 'd -> b h 1 d', h = self.heads,  b = b)
+
+        k = torch.cat((nk, k), dim = -2)
+        v = torch.cat((nv, v), dim = -2)
+
+        q = q * self.scale
+
+        # cosine sim attention
+
+        if self.cosine_sim_attn:
+            q, k = map(l2norm, (q, k))
+
+        # similarities
+
+        sim = einsum('b h i d, b h j d -> b h i j', q, k) * self.cosine_sim_scale
+
+        # masking
+
+        max_neg_value = -torch.finfo(sim.dtype).max
+
+        if exists(mask):
+            mask = F.pad(mask, (1, 0), value = True)
+            mask = rearrange(mask, 'b j -> b 1 1 j')
+            sim = sim.masked_fill(~mask, max_neg_value)
+
+        attn = sim.softmax(dim = -1, dtype = torch.float32)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head = 64,
+        heads = 8,
+        cosine_sim_attn = False
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5 if not cosine_sim_attn else 1
+        self.cosine_sim_attn = cosine_sim_attn
+        self.cosine_sim_scale = 16 if cosine_sim_attn else 1
+
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm = nn.LayerNorm(dim)
+        self.norm_latents = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim, bias = False),
+            nn.LayerNorm(dim)
+        )
+
+    def forward(self, x, latents, mask = None):
+        x = self.norm(x)
+        latents = self.norm_latents(latents)
+
+        b, h = x.shape[0], self.heads
+
+        q = self.to_q(latents)
+
+        # the paper differs from Perceiver in which they also concat the key / values derived from the latents to be attended to
+        kv_input = torch.cat((x, latents), dim = -2)
+        k, v = self.to_kv(kv_input).chunk(2, dim = -1)
+
+        q, k, v = rearrange_many((q, k, v), 'b n (h d) -> b h n d', h = h)
+
+        q = q * self.scale
+
+        # cosine sim attention
+
+        if self.cosine_sim_attn:
+            q, k = map(l2norm, (q, k))
+
+        # similarities and masking
+
+        sim = einsum('... i d, ... j d  -> ... i j', q, k) * self.cosine_sim_scale
+
+        if exists(mask):
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = F.pad(mask, (0, latents.shape[-2]), value = True)
+            mask = rearrange(mask, 'b j -> b 1 1 j')
+            sim = sim.masked_fill(~mask, max_neg_value)
+
+        # attention
+
+        attn = sim.softmax(dim = -1)
+
+        out = einsum('... i j, ... j d -> ... i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)', h = h)
+        return self.to_out(out)
+
+
+def FeedForward(dim, mult = 2):
+    hidden_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, hidden_dim, bias = False),
+        nn.GELU(),
+        nn.LayerNorm(hidden_dim),
+        nn.Linear(hidden_dim, dim, bias = False)
+    )
+
+def exists(val):
+    return val is not None
+
+
+def masked_mean(t, *, dim, mask = None):
+    if not exists(mask):
+        return t.mean(dim = dim)
+    denom = mask.sum(dim = dim, keepdim = True)
+    mask = rearrange(mask, 'b n -> b n 1')
+    masked_t = t.masked_fill(~mask, 0.)
+
+    return masked_t.sum(dim = dim) / denom.clamp(min = 1e-5)
+
+class PerceiverResampler(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head = 64,
+        heads = 8,
+        num_latents = 64,
+        num_latents_mean_pooled = 4, # number of latents derived from mean pooled representation of the sequence
+        max_seq_len = 512,
+        ff_mult = 4,
+        cosine_sim_attn = False
+    ):
+        super().__init__()
+        self.pos_emb = nn.Embedding(max_seq_len, dim)
+
+        self.latents = nn.Parameter(torch.randn(num_latents, dim))
+
+        self.to_latents_from_mean_pooled_seq = None
+
+        if num_latents_mean_pooled > 0:
+            self.to_latents_from_mean_pooled_seq = nn.Sequential(
+                nn.LayerNorm(dim),
+                nn.Linear(dim, dim * num_latents_mean_pooled),
+                Rearrange('b (n d) -> b n d', n = num_latents_mean_pooled)
+            )
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                PerceiverAttention(dim = dim, dim_head = dim_head, heads = heads, cosine_sim_attn = cosine_sim_attn),
+                FeedForward(dim = dim, mult = ff_mult)
+            ]))
+
+    def forward(self, x, mask = None):
+        n, device = x.shape[1], x.device
+        pos_emb = self.pos_emb(torch.arange(n, device = device))
+
+        x_with_pos = x + pos_emb
+
+        latents = repeat(self.latents, 'n d -> b n d', b = x.shape[0])
+
+        if exists(self.to_latents_from_mean_pooled_seq):
+            meanpooled_seq = masked_mean(x, dim = 1, mask = torch.ones(x.shape[:2], device = x.device, dtype = torch.bool))
+            meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq)
+            latents = torch.cat((meanpooled_latents, latents), dim = -2)
+
+        for attn, ff in self.layers:
+            latents = attn(x_with_pos, latents, mask = mask) + latents
+            latents = ff(latents) + latents
+
+        return latents
+
+
+
+
 class Upsample(nn.Module):
   def __init__(self, channels, with_conv=False):
     super().__init__()
@@ -616,4 +886,4 @@ class ResnetBlockDDPM(nn.Module):
         x = self.Conv_2(x)
       else:
         x = self.NIN_0(x)
-    return x + h
+    return x + h

score_sde/models/ncsnpp_generator_adagn.py

@@ -66,8 +66,10 @@ class NCSNpp(nn.Module):
     self.not_use_tanh = config.not_use_tanh
     self.act = act = nn.SiLU()
     self.z_emb_dim = z_emb_dim = config.z_emb_dim
-    
     self.nf = nf = config.num_channels_dae
+    self.cond_proj = nn.Linear(config.cond_size, self.nf*4) 
+    self.cond_proj.weight.data = default_initializer()(self.cond_proj.weight.shape)
+
     ch_mult = config.ch_mult
     self.num_res_blocks = num_res_blocks = config.num_res_blocks
     self.attn_resolutions = attn_resolutions = config.attn_resolutions
@@ -115,10 +117,12 @@ class NCSNpp(nn.Module):
       modules.append(nn.Linear(nf * 4, nf * 4))
       modules[-1].weight.data = default_initializer()(modules[-1].weight.shape)
       nn.init.zeros_(modules[-1].bias)
-
-    AttnBlock = functools.partial(layerspp.AttnBlockpp,
-                                  init_scale=init_scale,
-                                  skip_rescale=skip_rescale)
+    if config.cross_attention:
+      AttnBlock = functools.partial(layers.CondAttnBlock, context_dim=config.cond_size)
+    else:
+      AttnBlock = functools.partial(layerspp.AttnBlockpp,
+                                    init_scale=init_scale,
+                                    skip_rescale=skip_rescale)
 
     Upsample = functools.partial(layerspp.Upsample,
                                  with_conv=resamp_with_conv, fir=fir, fir_kernel=fir_kernel)
@@ -277,7 +281,7 @@ class NCSNpp(nn.Module):
     self.z_transform = nn.Sequential(*mapping_layers)
     
 
-  def forward(self, x, time_cond, z):
+  def forward(self, x, time_cond, z, cond=None):
     # timestep/noise_level embedding; only for continuous training
     zemb = self.z_transform(z)
     modules = self.all_modules
@@ -296,9 +300,14 @@ class NCSNpp(nn.Module):
 
     else:
       raise ValueError(f'embedding type {self.embedding_type} unknown.')
-
+    
+    if cond is not None:
+      cond_pooled, cond, cond_mask = cond
+    
     if self.conditional:
       temb = modules[m_idx](temb)
+      if cond is not None:
+        temb = temb + self.cond_proj(cond_pooled)
       m_idx += 1
       temb = modules[m_idx](self.act(temb))
       m_idx += 1
@@ -322,7 +331,10 @@ class NCSNpp(nn.Module):
         h = modules[m_idx](hs[-1], temb, zemb)
         m_idx += 1
         if h.shape[-1] in self.attn_resolutions:
-          h = modules[m_idx](h)
+          if type(modules[m_idx]) == layers.CondAttnBlock:
+            h = modules[m_idx](h, cond, cond_mask)
+          else:
+            h = modules[m_idx](h)
           m_idx += 1
 
         hs.append(h)
@@ -354,7 +366,10 @@ class NCSNpp(nn.Module):
     h = hs[-1]
     h = modules[m_idx](h, temb, zemb)
     m_idx += 1
-    h = modules[m_idx](h)
+    if type(modules[m_idx]) == layers.CondAttnBlock:
+      h = modules[m_idx](h, cond, cond_mask)
+    else:
+      h = modules[m_idx](h)
     m_idx += 1
     h = modules[m_idx](h, temb, zemb)
     m_idx += 1
@@ -368,7 +383,10 @@ class NCSNpp(nn.Module):
         m_idx += 1
 
       if h.shape[-1] in self.attn_resolutions:
-        h = modules[m_idx](h)
+        if type(modules[m_idx]) == layers.CondAttnBlock:
+          h = modules[m_idx](h, cond, cond_mask)
+        else:
+          h = modules[m_idx](h)
         m_idx += 1
 
       if self.progressive != 'none':
@@ -429,3 +447,4 @@ class NCSNpp(nn.Module):
         return torch.tanh(h)
     else:
         return h
+

scripts/init.sh

@@ -0,0 +1,34 @@
+machine=$(cat /etc/FZJ/systemname)
+if [[ "$machine" == jurecadc ]]; 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 /p/project/covidnetx/environments/jureca_2022/bin/activate 
+fi
+if [[ "$machine" == juwelsbooster ]]; 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 /p/project/covidnetx/environments/juwels_booster_2022/bin/activate 
+fi
+if [[ "$machine" == jusuf ]]; then
+    echo not supported
+fi

scripts/init_2020.sh

@@ -0,0 +1,69 @@
+machine=$(cat /etc/FZJ/systemname)
+if [[ "$machine" == jurecadc ]]; then
+    export CUDA_VISIBLE_DEVICES=0,1,2,3
+    #ml use $OTHERSTAGES
+    #ml Stages/2020
+    #ml GCC/9.3.0
+    #ml OpenMPI/4.1.0rc1
+    #ml CUDA/11.0
+    #ml cuDNN/8.0.2.39-CUDA-11.0
+    #ml NCCL/2.8.3-1-CUDA-11.0
+    #ml PyTorch
+    #ml Horovod/0.20.3-Python-3.8.5
+    #ml scikit
+    #source /p/project/covidnetx/environments/jureca/bin/activate 
+    ml purge
+    ml use $OTHERSTAGES
+    ml Stages/2020
+    ml GCC/10.3.0
+    ml OpenMPI/4.1.1
+    ml Horovod/0.23.0-Python-3.8.5
+    ml scikit
+    source /p/project/covidnetx/environments/jureca/bin/activate 
+fi
+if [[ "$machine" == juwelsbooster ]]; then
+    export CUDA_VISIBLE_DEVICES=0,1,2,3
+    #ml use $OTHERSTAGES
+    #ml Stages/2020
+    #ml GCC/9.3.0
+    #ml OpenMPI/4.1.0rc1
+    #ml CUDA/11.0
+    #ml cuDNN/8.0.2.39-CUDA-11.0
+    #ml NCCL/2.8.3-1-CUDA-11.0
+    #ml PyTorch
+    #ml Horovod/0.20.3-Python-3.8.5
+    #ml scikit
+
+    #ml Stages/2021
+    #ml GCC
+    #ml OpenMPI
+    #ml CUDA
+    #ml cuDNN
+    #ml NCCL
+    #ml PyTorch
+    #ml Horovod
+    #ml scikit
+
+    ml purge
+    ml use $OTHERSTAGES
+    ml Stages/2020
+    ml GCC/10.3.0
+    ml OpenMPI/4.1.1
+    ml Horovod/0.23.0-Python-3.8.5
+    ml scikit
+    source /p/project/covidnetx/environments/juwels_booster/bin/activate 
+fi
+if [[ "$machine" == jusuf ]]; then
+    ml purge
+    ml use $OTHERSTAGES
+    ml Stages/2020
+    ml GCC/9.3.0
+    ml OpenMPI/4.1.0rc1
+    ml CUDA/11.0
+    ml cuDNN/8.0.2.39-CUDA-11.0
+    ml NCCL/2.8.3-1-CUDA-11.0
+    ml PyTorch
+    ml Horovod/0.20.3-Python-3.8.5
+    #ml scikit
+    source /p/project/covidnetx/environments/jusuf/bin/activate
+fi

scripts/init_2022.sh

@@ -0,0 +1,34 @@
+machine=$(cat /etc/FZJ/systemname)
+if [[ "$machine" == jurecadc ]]; 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 /p/project/covidnetx/environments/jureca_2022/bin/activate 
+fi
+if [[ "$machine" == juwelsbooster ]]; 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 /p/project/covidnetx/environments/juwels_booster_2022/bin/activate 
+fi
+if [[ "$machine" == jusuf ]]; then
+    echo not supported
+fi

scripts/run_jurecadc_ddp.sh

@@ -0,0 +1,17 @@
+#!/bin/bash -x
+#SBATCH --account=zam
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=4
+#SBATCH --cpus-per-task=24
+#SBATCH --time=06:00:00
+#SBATCH --gres=gpu:4
+#SBATCH --partition=dc-gpu
+source set_torch_distributed_vars.sh
+#source scripts/init_2022.sh
+#source scripts/init_2020.sh
+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
+srun python -u $*

scripts/run_jusuf_ddp.sh

@@ -0,0 +1,14 @@
+#!/bin/bash -x
+#SBATCH --account=zam
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=24
+#SBATCH --time=06:00:00
+#SBATCH --gres=gpu:1
+#SBATCH --partition=gpus
+source set_torch_distributed_vars.sh
+source scripts/init.sh
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+echo "Job id: $SLURM_JOB_ID"
+export TOKENIZERS_PARALLELISM=false
+srun python -u $*

scripts/run_juwelsbooster_ddp.sh

@@ -0,0 +1,17 @@
+#!/bin/bash -x
+#SBATCH --account=covidnetx
+#SBATCH --nodes=4
+#SBATCH --ntasks-per-node=4
+#SBATCH --cpus-per-task=24
+#SBATCH --time=06:00:00
+#SBATCH --gres=gpu:4
+#SBATCH --partition=booster
+source set_torch_distributed_vars.sh
+#source scripts/init_2022.sh
+#source scripts/init_2020.sh
+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
+srun python -u $*

t5.py

@@ -0,0 +1,99 @@
+import torch
+import transformers
+from transformers import T5Tokenizer, T5EncoderModel, T5Config
+
+transformers.logging.set_verbosity_error()
+
+def exists(val):
+    return val is not None
+
+# config
+
+MAX_LENGTH = 256
+
+DEFAULT_T5_NAME = 'google/t5-v1_1-base'
+
+T5_CONFIGS = {}
+
+# singleton globals
+
+def get_tokenizer(name):
+    tokenizer = T5Tokenizer.from_pretrained(name)
+    return tokenizer
+
+def get_model(name):
+    model = T5EncoderModel.from_pretrained(name)
+    return model
+
+def get_model_and_tokenizer(name):
+    global T5_CONFIGS
+
+    if name not in T5_CONFIGS:
+        T5_CONFIGS[name] = dict()
+    if "model" not in T5_CONFIGS[name]:
+        T5_CONFIGS[name]["model"] = get_model(name)
+    if "tokenizer" not in T5_CONFIGS[name]:
+        T5_CONFIGS[name]["tokenizer"] = get_tokenizer(name)
+
+    return T5_CONFIGS[name]['model'], T5_CONFIGS[name]['tokenizer']
+
+def get_encoded_dim(name):
+    if name not in T5_CONFIGS:
+        # avoids loading the model if we only want to get the dim
+        config = T5Config.from_pretrained(name)
+        T5_CONFIGS[name] = dict(config=config)
+    elif "config" in T5_CONFIGS[name]:
+        config = T5_CONFIGS[name]["config"]
+    elif "model" in T5_CONFIGS[name]:
+        config = T5_CONFIGS[name]["model"].config
+    else:
+        assert False
+    return config.d_model
+
+class T5Encoder(torch.nn.Module):
+
+    def __init__(self, name=DEFAULT_T5_NAME, max_length=MAX_LENGTH, padding='longest', masked_mean=False):
+        super().__init__()
+        self.name = name
+        self.t5, self.tokenizer = get_model_and_tokenizer(name)
+        self.max_length = max_length
+        self.output_size = get_encoded_dim(name)
+        self.padding = padding
+        self.masked_mean = masked_mean
+
+    def forward(self, x, return_only_pooled=True):
+        encoded = self.tokenizer.batch_encode_plus(
+            x,
+            return_tensors = "pt",
+            padding = self.padding,
+            max_length = self.max_length,
+            truncation = True
+        )
+        device = next(self.t5.parameters()).device
+        input_ids = encoded.input_ids.to(device)
+        attn_mask = encoded.attention_mask.to(device).bool()
+        output = self.t5(input_ids = input_ids, attention_mask = attn_mask)
+        encoded_text = output.last_hidden_state.detach()
+        # return encoded_text[:, 0]
+        # print(input_ids)
+        # print(attn_mask)
+        #if self.masked_mean:
+        pooled =  masked_mean(encoded_text, dim=1, mask=attn_mask)
+        if return_only_pooled:
+            return pooled
+        else:
+            return pooled, encoded_text, attn_mask
+        #else:
+        #    return encoded_text.mean(dim=1)
+
+
+from einops import rearrange
+def masked_mean(t, *, dim, mask = None):
+    if not exists(mask):
+        return t.mean(dim = dim)
+
+    denom = mask.sum(dim = dim, keepdim = True)
+    mask = rearrange(mask, 'b n -> b n 1')
+    masked_t = t.masked_fill(~mask, 0.)
+
+    return masked_t.sum(dim = dim) / denom.clamp(min = 1e-5)

test_ddgan.py

@@ -7,12 +7,12 @@
 import argparse
 import torch
 import numpy as np
-
+import time
 import os
-
+import json
 import torchvision
 from score_sde.models.ncsnpp_generator_adagn import NCSNpp
-from pytorch_fid.fid_score import calculate_fid_given_paths
+import t5
 
 #%% Diffusion coefficients 
 def var_func_vp(t, beta_min, beta_max):
@@ -112,7 +112,7 @@ def sample_posterior(coefficients, x_0,x_t, t):
     
     return sample_x_pos
 
-def sample_from_model(coefficients, generator, n_time, x_init, T, opt):
+def sample_from_model(coefficients, generator, n_time, x_init, T, opt, cond=None):
     x = x_init
     with torch.no_grad():
         for i in reversed(range(n_time)):
@@ -120,17 +120,70 @@ def sample_from_model(coefficients, generator, n_time, x_init, T, opt):
             
             t_time = t
             latent_z = torch.randn(x.size(0), opt.nz, device=x.device)#.to(x.device)
-            x_0 = generator(x, t_time, latent_z)
+            x_0 = generator(x, t_time, latent_z, cond=cond)
             x_new = sample_posterior(coefficients, x_0, x, t)
             x = x_new.detach()
         
     return x
 
+
+def sample_from_model_classifier_free_guidance(coefficients, generator, n_time, x_init, T, opt, text_encoder, cond=None, guidance_scale=0):
+    x = x_init
+    null = text_encoder([""] * len(x_init), return_only_pooled=False)
+    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)
+            
+            x_0_uncond = generator(x, t_time, latent_z, cond=null)
+            x_0_cond = generator(x, t_time, latent_z, cond=cond)
+
+            eps_uncond = (x - torch.sqrt(coefficients.alphas_cumprod[i]) * x_0_uncond) / torch.sqrt(1 - coefficients.alphas_cumprod[i])
+            eps_cond = (x - torch.sqrt(coefficients.alphas_cumprod[i]) * x_0_cond) / torch.sqrt(1 - coefficients.alphas_cumprod[i])
+            
+            # 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)
+           
+
+            # Dynamic thresholding
+            q = args.dynamic_thresholding_percentile
+            print("Before", x_0.min(), x_0.max())
+            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)
+            print("After", x_0.min(), x_0.max())
+            
+            x_new = sample_posterior(coefficients, x_0, x, t)
+            
+            # Dynamic thresholding
+            # q = args.dynamic_thresholding_percentile
+            # shape = x_new.shape
+            # x_new_v = x_new.view(shape[0], -1)
+            # d = torch.quantile(torch.abs(x_new_v), q, dim=1, keepdim=True)
+            # d = torch.maximum(d, torch.ones_like(d))
+            # d.clamp_(min = 1.)
+            # x_new_v = torch.clamp(x_new_v, -d, d) / d
+            # x_new = x_new_v.view(shape)
+            x = x_new.detach()
+        
+    return x
+
+
 #%%
 def sample_and_test(args):
-    torch.manual_seed(42)
+    torch.manual_seed(args.seed)
     device = 'cuda:0'
-    
+    text_encoder = t5.T5Encoder(name=args.text_encoder, masked_mean=args.masked_mean).to(device)
+    args.cond_size = text_encoder.output_size
+    # cond = text_encoder([str(yi%10) for yi in range(args.batch_size)])
+
     if args.dataset == 'cifar10':
         real_img_dir = 'pytorch_fid/cifar10_train_stat.npy'
     elif args.dataset == 'celeba_256':
@@ -157,7 +210,6 @@ def sample_and_test(args):
     
     pos_coeff = Posterior_Coefficients(args, device)
         
-    iters_needed = 50000 //args.batch_size
     
     save_dir = "./generated_samples/{}".format(args.dataset)
     
@@ -165,25 +217,90 @@ def sample_and_test(args):
         os.makedirs(save_dir)
     
     if args.compute_fid:
-        for i in range(iters_needed):
+        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
+
+        texts = open(args.cond_text).readlines()
+        #iters_needed = len(texts) // args.batch_size
+        #texts = list(map(lambda s:s.strip(), texts))
+        #ntimes = max(30000 // len(texts), 1)
+        #texts = texts * ntimes
+        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 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 = []
+        for b in range(0, len(texts), args.batch_size):
+            text = texts[b:b+args.batch_size]
             with torch.no_grad():
-                x_t_1 = torch.randn(args.batch_size, args.num_channels,args.image_size, args.image_size).to(device)
-                fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args)
-                
+                cond = text_encoder(text, return_only_pooled=False)
+                bs = len(text)
+                t0 = time.time()
+                x_t_1 = torch.randn(bs, args.num_channels,args.image_size, args.image_size).to(device)
+                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:
+                    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))
-                print('generating batch ', i)
-        
-        paths = [save_dir, real_img_dir]
-    
-        kwargs = {'batch_size': 100, 'device': device, 'dims': 2048}
-        fid = calculate_fid_given_paths(paths=paths, **kwargs)
+                """
+                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 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)
+                """
+            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/{}/{}/fid_{}.json'.format(args.dataset, args.exp, args.epoch_id)
+        results = {
+            "fid": fid,
+        }
+        results.update(vars(args))
+        with open(dest, "w") as fd:
+            json.dump(results, fd)
         print('FID = {}'.format(fid))
     else:
+        cond = text_encoder([args.cond_text] * args.batch_size, return_only_pooled=False)
         x_t_1 = torch.randn(args.batch_size, args.num_channels,args.image_size, args.image_size).to(device)
-        fake_sample = sample_from_model(pos_coeff, netG, args.num_timesteps, x_t_1,T,  args)
+        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:
+            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)
         torchvision.utils.save_image(fake_sample, './samples_{}.jpg'.format(args.dataset))
 
@@ -198,6 +315,13 @@ if __name__ == '__main__':
     parser.add_argument('--compute_fid', action='store_true', default=False,
                             help='whether or not compute FID')
     parser.add_argument('--epoch_id', type=int,default=1000)
+    parser.add_argument('--guidance_scale', type=float,default=0)
+    parser.add_argument('--dynamic_thresholding_percentile', type=float,default=0)
+    parser.add_argument('--cond_text', type=str,default="0")
+
+    parser.add_argument('--cross_attention', action='store_true',default=False)
+
+    
     parser.add_argument('--num_channels', type=int, default=3,
                             help='channel of image')
     parser.add_argument('--centered', action='store_false', default=True,
@@ -262,6 +386,8 @@ if __name__ == '__main__':
     parser.add_argument('--z_emb_dim', type=int, default=256)
     parser.add_argument('--t_emb_dim', type=int, default=256)
     parser.add_argument('--batch_size', type=int, default=200, help='sample generating batch size')
+    parser.add_argument('--text_encoder', type=str, default="google/t5-v1_1-base")
+    parser.add_argument('--masked_mean', action='store_true',default=False)
         
 
 
@@ -272,4 +398,4 @@ if __name__ == '__main__':
     sample_and_test(args)
     
    
-                
+                

train_ddgan.py

@@ -18,7 +18,7 @@ import torch.optim as optim
 import torchvision
 
 import torchvision.transforms as transforms
-from torchvision.datasets import CIFAR10
+from torchvision.datasets import CIFAR10, ImageFolder
 from datasets_prep.lsun import LSUN
 from datasets_prep.stackmnist_data import StackedMNIST, _data_transforms_stacked_mnist
 from datasets_prep.lmdb_datasets import LMDBDataset
@@ -27,6 +27,11 @@ from datasets_prep.lmdb_datasets import LMDBDataset
 from torch.multiprocessing import Process
 import torch.distributed as dist
 import shutil
+import logging
+import t5
+def log_and_continue(exn):
+    logging.warning(f'Handling webdataset error ({repr(exn)}). Ignoring.')
+    return True
 
 def copy_source(file, output_dir):
     shutil.copyfile(file, os.path.join(output_dir, os.path.basename(file)))
@@ -172,7 +177,7 @@ def sample_posterior(coefficients, x_0,x_t, t):
     
     return sample_x_pos
 
-def sample_from_model(coefficients, generator, n_time, x_init, T, opt):
+def sample_from_model(coefficients, generator, n_time, x_init, T, opt, cond=None):
     x = x_init
     with torch.no_grad():
         for i in reversed(range(n_time)):
@@ -180,13 +185,15 @@ def sample_from_model(coefficients, generator, n_time, x_init, T, opt):
           
             t_time = t
             latent_z = torch.randn(x.size(0), opt.nz, device=x.device)
-            x_0 = generator(x, t_time, latent_z)
+            x_0 = generator(x, t_time, latent_z, cond=cond)
             x_new = sample_posterior(coefficients, x_0, x, t)
             x = x_new.detach()
         
     return x
 
-#%%
+
+from utils import ResampledShards2
+
 def train(rank, gpu, args):
     from score_sde.models.discriminator import Discriminator_small, Discriminator_large
     from score_sde.models.ncsnpp_generator_adagn import NCSNpp
@@ -236,37 +243,81 @@ def train(rank, gpu, args):
                 transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
             ])
         dataset = LMDBDataset(root='/datasets/celeba-lmdb/', name='celeba', train=True, transform=train_transform)
-      
-    
+    elif args.dataset == "image_folder":
+        train_transform = transforms.Compose([
+                transforms.Resize(args.image_size),
+                transforms.CenterCrop(args.image_size),
+                # transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+            ])
+        dataset = ImageFolder(root=args.dataset_root, transform=train_transform)
+    elif args.dataset == 'wds':
+        import webdataset as wds
+        train_transform = transforms.Compose([
+                transforms.Resize(args.image_size),
+                transforms.CenterCrop(args.image_size),
+                # transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+            ])
+        # pipeline = [wds.SimpleShardList(args.dataset_root)]
+        pipeline = [ResampledShards2(args.dataset_root)]
+        pipeline.extend([
+            wds.split_by_node,
+            wds.split_by_worker,
+            wds.tarfile_to_samples(handler=log_and_continue),
+        ])
+        pipeline.extend([
+            wds.decode("pilrgb", handler=log_and_continue),
+            wds.rename(image="jpg;png"),
+            wds.map_dict(image=train_transform),
+            wds.to_tuple("image","txt"),
+            wds.batched(batch_size, partial=False),
+        ])
+        dataset = wds.DataPipeline(*pipeline)
+        data_loader = wds.WebLoader(
+            dataset,
+            batch_size=None,
+            shuffle=False,
+            num_workers=8,
+        )
     
-    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset,
-                                                                    num_replicas=args.world_size,
-                                                                    rank=rank)
-    data_loader = torch.utils.data.DataLoader(dataset,
-                                               batch_size=batch_size,
+    if args.dataset != "wds":
+        train_sampler = torch.utils.data.distributed.DistributedSampler(dataset,
+                                                                        num_replicas=args.world_size,
+                                                                        rank=rank)
+        data_loader = torch.utils.data.DataLoader(dataset,
+                                                   batch_size=batch_size,
                                                shuffle=False,
                                                num_workers=4,
+                                               drop_last=True,
                                                pin_memory=True,
-                                               sampler=train_sampler,
-                                               drop_last = True)
-    
+                                               sampler=train_sampler,)
+    text_encoder = t5.T5Encoder(name=args.text_encoder, masked_mean=args.masked_mean).to(device)
+    args.cond_size = text_encoder.output_size
     netG = NCSNpp(args).to(device)
+    nb_params = 0
+    for param in netG.parameters():
+        nb_params += param.flatten().shape[0]
+    print("Number of generator parameters:", nb_params)
     
 
     if args.dataset == 'cifar10' or args.dataset == 'stackmnist':    
         netD = Discriminator_small(nc = 2*args.num_channels, ngf = args.ngf,
                                t_emb_dim = args.t_emb_dim,
+                               cond_size=text_encoder.output_size,
                                act=nn.LeakyReLU(0.2)).to(device)
     else:
         netD = Discriminator_large(nc = 2*args.num_channels, ngf = args.ngf, 
                                    t_emb_dim = args.t_emb_dim,
+                                cond_size=text_encoder.output_size,
                                    act=nn.LeakyReLU(0.2)).to(device)
     
     broadcast_params(netG.parameters())
     broadcast_params(netD.parameters())
     
     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))
     
     if args.use_ema:
@@ -297,9 +348,9 @@ def train(rank, gpu, args):
     pos_coeff = Posterior_Coefficients(args, device)
     T = get_time_schedule(args, device)
     
-    if args.resume:
-        checkpoint_file = os.path.join(exp_path, 'content.pth')
-        checkpoint = torch.load(checkpoint_file, map_location=device)
+    checkpoint_file = os.path.join(exp_path, 'content.pth')
+    if args.resume and os.path.exists(checkpoint_file):
+        checkpoint = torch.load(checkpoint_file, map_location="cpu")
         init_epoch = checkpoint['epoch']
         epoch = init_epoch
         netG.load_state_dict(checkpoint['netG_dict'])
@@ -319,9 +370,22 @@ def train(rank, gpu, args):
     
     
     for epoch in range(init_epoch, args.num_epoch+1):
-        train_sampler.set_epoch(epoch)
+        if args.dataset == "wds":
+            os.environ["WDS_EPOCH"] = str(epoch)
+        else:
+            train_sampler.set_epoch(epoch)
        
         for iteration, (x, y) in enumerate(data_loader):
+            if args.dataset != "wds":
+                y = [str(yi) for yi in y.tolist()]
+            
+            if args.classifier_free_guidance_proba:
+                u = (np.random.uniform(size=len(y)) <= args.classifier_free_guidance_proba).tolist()
+                y = ["" if ui else yi for yi,ui in zip(y, u)]
+
+            with torch.no_grad():
+                cond_pooled, cond, cond_mask = text_encoder(y, return_only_pooled=False)
+
             for p in netD.parameters():  
                 p.requires_grad = True  
         
@@ -339,7 +403,7 @@ def train(rank, gpu, args):
             
     
             # train with real
-            D_real = netD(x_t, t, x_tp1.detach()).view(-1)
+            D_real = netD(x_t, t, x_tp1.detach(), cond=cond_pooled).view(-1)
             
             errD_real = F.softplus(-D_real)
             errD_real = errD_real.mean()
@@ -375,10 +439,10 @@ def train(rank, gpu, args):
             latent_z = torch.randn(batch_size, nz, device=device)
             
          
-            x_0_predict = netG(x_tp1.detach(), t, latent_z)
+            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()).view(-1)
+            output = netD(x_pos_sample, t, x_tp1.detach(), cond=cond_pooled).view(-1)
                 
             
             errD_fake = F.softplus(output)
@@ -407,11 +471,10 @@ def train(rank, gpu, args):
             
             
                 
-           
-            x_0_predict = netG(x_tp1.detach(), t, latent_z)
+            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()).view(-1)
+            output = netD(x_pos_sample, t, x_tp1.detach(), cond=cond_pooled).view(-1)
                
             
             errG = F.softplus(-output)
@@ -426,7 +489,27 @@ def train(rank, gpu, args):
             if iteration % 100 == 0:
                 if rank == 0:
                     print('epoch {} iteration{}, G Loss: {}, D Loss: {}'.format(epoch,iteration, errG.item(), errD.item()))
-        
+            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))
+                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.')
+                        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()
@@ -437,7 +520,7 @@ def train(rank, gpu, args):
                 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)
+            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:
@@ -449,6 +532,7 @@ 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'))
                 
             if epoch % args.save_ckpt_every == 0:
                 if args.use_ema:
@@ -462,11 +546,19 @@ def train(rank, gpu, args):
 
 def init_processes(rank, size, fn, args):
     """ Initialize the distributed environment. """
+
+    import os
+
+    args.rank = int(os.environ['SLURM_PROCID'])
+    args.world_size =  int(os.getenv("SLURM_NTASKS"))
+    args.local_rank = int(os.environ['SLURM_LOCALID'])
+    print(args.rank, args.world_size)
+    args.master_address = os.getenv("SLURM_LAUNCH_NODE_IPADDR")
     os.environ['MASTER_ADDR'] = args.master_address
-    os.environ['MASTER_PORT'] = '6020'
+    os.environ['MASTER_PORT'] = "12345"
     torch.cuda.set_device(args.local_rank)
     gpu = args.local_rank
-    dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=size)
+    dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=args.world_size)
     fn(rank, gpu, args)
     dist.barrier()
     cleanup()  
@@ -480,7 +572,10 @@ 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('--text_encoder', type=str, default="google/t5-v1_1-base")
+    parser.add_argument('--cross_attention', action='store_true',default=False)
+
     parser.add_argument('--image_size', type=int, default=32,
                             help='size of image')
     parser.add_argument('--num_channels', type=int, default=3,
@@ -492,7 +587,7 @@ if __name__ == '__main__':
                             help='beta_min for diffusion')
     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')
@@ -534,6 +629,7 @@ if __name__ == '__main__':
     #geenrator and training
     parser.add_argument('--exp', default='experiment_cifar_default', help='name of experiment')
     parser.add_argument('--dataset', default='cifar10', help='name of dataset')
+    parser.add_argument('--dataset_root', default='', help='name of dataset')
     parser.add_argument('--nz', type=int, default=100)
     parser.add_argument('--num_timesteps', type=int, default=4)
 
@@ -577,26 +673,28 @@ if __name__ == '__main__':
 
    
     args = parser.parse_args()
-    args.world_size = args.num_proc_node * args.num_process_per_node
-    size = args.num_process_per_node
-
-    if size > 1:
-        processes = []
-        for rank in range(size):
-            args.local_rank = rank
-            global_rank = rank + args.node_rank * args.num_process_per_node
-            global_size = args.num_proc_node * args.num_process_per_node
-            args.global_rank = global_rank
-            print('Node rank %d, local proc %d, global proc %d' % (args.node_rank, rank, global_rank))
-            p = Process(target=init_processes, args=(global_rank, global_size, train, args))
-            p.start()
-            processes.append(p)
-            
-        for p in processes:
-            p.join()
-    else:
-        print('starting in debug mode')
+    # args.world_size = args.num_proc_node * args.num_process_per_node
+    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)
+    # if size > 1:
+        # processes = []
+        # for rank in range(size):
+            # args.local_rank = rank
+            # global_rank = rank + args.node_rank * args.num_process_per_node
+            # global_size = args.num_proc_node * args.num_process_per_node
+            # args.global_rank = global_rank
+            # print('Node rank %d, local proc %d, global proc %d' % (args.node_rank, rank, global_rank))
+            # p = Process(target=init_processes, args=(global_rank, global_size, train, args))
+            # p.start()
+            # processes.append(p)
+            
+        # for p in processes:
+            # p.join()
+    # else:
+        # print('starting in debug mode')
         
-        init_processes(0, size, train, args)
+        # init_processes(0, size, train, args)
    
-                
+                

utils.py

@@ -0,0 +1,67 @@
+from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler, IterableDataset, get_worker_info
+import braceexpand
+import random
+import sys
+def pytorch_worker_seed():
+    """get dataloader worker seed from pytorch"""
+    worker_info = get_worker_info()
+    if worker_info is not None:
+        # favour the seed already created for pytorch dataloader workers if it exists
+        return worker_info.seed
+    # fallback to wds rank based seed
+    return wds.utils.pytorch_worker_seed()
+
+
+class SharedEpoch:
+    def __init__(self, epoch: int = 0):
+        self.shared_epoch = Value('i', epoch)
+
+    def set_value(self, epoch):
+        self.shared_epoch.value = epoch
+
+    def get_value(self):
+        return self.shared_epoch.value
+
+
+
+class ResampledShards2(IterableDataset):
+    """An iterable dataset yielding a list of urls."""
+
+    def __init__(
+        self,
+        urls,
+        nshards=sys.maxsize,
+        worker_seed=None,
+        deterministic=False,
+        epoch=-1,
+    ):
+        """Sample shards from the shard list with replacement.
+
+        :param urls: a list of URLs as a Python list or brace notation string
+        """
+        super().__init__()
+        #urls = wds.shardlists.expand_urls(urls)
+        urls = list(braceexpand.braceexpand(urls))
+        self.urls = urls
+        assert isinstance(self.urls[0], str)
+        self.nshards = nshards
+        self.rng = random.Random()
+        self.worker_seed = pytorch_worker_seed if worker_seed is None else worker_seed
+        self.deterministic = deterministic
+        self.epoch = epoch
+
+    def __iter__(self):
+        """Return an iterator over the shards."""
+        if isinstance(self.epoch, SharedEpoch):
+            epoch = self.epoch.get_value()
+        else:
+            # NOTE: this is epoch tracking is problematic in a multiprocess (dataloader workers or train)
+            # situation as different workers may wrap at different times (or not at all).
+            self.epoch += 1
+            epoch = self.epoch
+        if self.deterministic:
+            # reset seed w/ epoch if deterministic, worker seed should be deterministic due to arg.seed
+            self.rng.seed(self.worker_seed() + epoch)
+        for _ in range(self.nshards):
+            yield dict(url=self.rng.choice(self.urls))
+