extract submodel script and segmentation stage example config

Author: pesser

configs/sflckr_cond_stage.yaml

@@ -0,0 +1,43 @@
+model:
+  base_learning_rate: 4.5e-06
+  target: taming.models.vqgan.VQSegmentationModel
+  params:
+    embed_dim: 256
+    n_embed: 1024
+    image_key: "segmentation"
+    n_labels: 182
+    ddconfig:
+      double_z: false
+      z_channels: 256
+      resolution: 256
+      in_channels: 182
+      out_ch: 182
+      ch: 128
+      ch_mult:
+      - 1
+      - 1
+      - 2
+      - 2
+      - 4
+      num_res_blocks: 2
+      attn_resolutions:
+      - 16
+      dropout: 0.0
+
+    lossconfig:
+      target: taming.modules.losses.segmentation.BCELossWithQuant
+      params:
+        codebook_weight: 1.0
+
+data:
+  target: cutlit.DataModuleFromConfig
+  params:
+    batch_size: 12
+    train:
+      target: taming.data.sflckr.Examples # adjust
+      params:
+        size: 256
+    validation:
+      target: taming.data.sflckr.Examples # adjust
+      params:
+        size: 256

scripts/extract_submodel.py

@@ -0,0 +1,17 @@
+import torch
+import sys
+
+if __name__ == "__main__":
+    inpath = sys.argv[1]
+    outpath = sys.argv[2]
+    submodel = "cond_stage_model"
+    if len(sys.argv) > 3:
+        submodel = sys.argv[3]
+
+    print("Extracting {} from {} to {}.".format(submodel, inpath, outpath))
+
+    sd = torch.load(inpath, map_location="cpu")
+    new_sd = {"state_dict": dict((k.split(".", 1)[-1],v)
+                                 for k,v in sd["state_dict"].items()
+                                 if k.startswith("cond_stage_model"))}
+    torch.save(new_sd, outpath)

taming/modules/losses/segmentation.py

@@ -0,0 +1,22 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BCELoss(nn.Module):
+    def forward(self, prediction, target):
+        loss = F.binary_cross_entropy_with_logits(prediction,target)
+        return loss, {}
+
+
+class BCELossWithQuant(nn.Module):
+    def __init__(self, codebook_weight=1.):
+        super().__init__()
+        self.codebook_weight = codebook_weight
+
+    def forward(self, qloss, target, prediction, split):
+        bce_loss = F.binary_cross_entropy_with_logits(prediction,target)
+        loss = bce_loss + self.codebook_weight*qloss
+        return loss, {"{}/total_loss".format(split): loss.clone().detach().mean(),
+                      "{}/bce_loss".format(split): bce_loss.detach().mean(),
+                      "{}/quant_loss".format(split): qloss.detach().mean()
+                      }