fix: allow non-scanned models (#168)

tools/train/config/medium/config.json

@@ -1,31 +0,0 @@
-{
-  "activation_dropout": 0.0,
-  "activation_function": "gelu",
-  "attention_dropout": 0.0,
-  "bos_token_id": 16385,
-  "d_model": 1408,
-  "decoder_attention_heads": 16,
-  "decoder_ffn_dim": 4096,
-  "decoder_layerdrop": 0.0,
-  "decoder_layers": 14,
-  "decoder_start_token_id": 16384,
-  "dropout": 0.0,
-  "encoder_attention_heads": 16,
-  "encoder_ffn_dim": 4096,
-  "encoder_layerdrop": 0.0,
-  "encoder_layers": 14,
-  "encoder_vocab_size": 50264,
-  "eos_token_id": 16385,
-  "gradient_checkpointing": false,
-  "image_length": 256,
-  "image_vocab_size": 16384,
-  "init_std": 0.01,
-  "is_encoder_decoder": true,
-  "max_text_length": 64,
-  "model_type": "dallebart",
-  "normalize_text": true,
-  "pad_token_id": 16385,
-  "scale_embedding": false,
-  "tie_word_embeddings": false,
-  "use_cache": true
-}

tools/train/train.py

@@ -536,6 +536,8 @@ def split_params(data):
             split["scanned_decoder"][k] = v
         else:
             split["standard"][k] = v
+    # remove empty keys
+    split = {k: v for k, v in split.items() if v}
     for k, v in split.items():
         split[k] = freeze(traverse_util.unflatten_dict(v))
     return split
@@ -544,7 +546,8 @@ def split_params(data):
 def unsplit_params(data):
     flat = {}
     for k in ["standard", "scanned_encoder", "scanned_decoder"]:
-        flat.update(traverse_util.flatten_dict(unfreeze(data[k])))
+        if k in data:
+            flat.update(traverse_util.flatten_dict(unfreeze(data[k])))
     return freeze(traverse_util.unflatten_dict(flat))
 
 
@@ -1483,7 +1486,7 @@ def main():
     logger.info("  Ready to start training")
     with mesh:
         for epoch in epochs:
-            state.replace(epoch=epoch)
+            state = state.replace(epoch=epoch)
             local_state["epoch"] = epoch
             # ======================== Training ================================
             metrics_logger.update_state_metrics(local_state)