allow passing a custom variable name mapping for preloading of parameters (#667)

Author: michelwi

returnn/tf/engine.py

@@ -1135,7 +1135,8 @@ def init_network_from_config(self, config=None, net_dict_post_proc=None):
           params_prefix=self_prefix, load_if_prefix=load_if_prefix,
           ignore_missing=opts.get("ignore_missing", False),
           ignore_params=opts.get("ignore_params", ()),
-          ignore_params_prefixes=opts.get("ignore_params_prefixes", ()))
+          ignore_params_prefixes=opts.get("ignore_params_prefixes", ()),
+          var_name_mapping=opts.get("var_name_mapping", {}))
         # `set_as_custom_init` is also a marker for the vars, that they are preloaded,
         # such that further checkpoint loaders will not load them again.
         loader.set_as_custom_init()

returnn/tf/network.py

@@ -3470,7 +3470,7 @@ class CustomCheckpointLoader:
   """
 
   def __init__(self, filename, saveable_params, params_prefix="", load_if_prefix="", ignore_missing=False,
-               ignore_params=(), ignore_params_prefixes=(),
+               ignore_params=(), ignore_params_prefixes=(), var_name_mapping=None,
                network=None):
     """
     :param str filename: filepattern for NewCheckpointReader
@@ -3482,13 +3482,16 @@ def __init__(self, filename, saveable_params, params_prefix="", load_if_prefix="
       however, if there is no single var in the checkpoint, this is still an error.
     :param typing.Container[str] ignore_params: these param (by name) will not be loaded
     :param typing.Iterable[str] ignore_params_prefixes: these param (by prefix name) will not be loaded
+    :param dict[str,str] var_name_mapping: defines a custom mapping (new_name -> name_in_checkpoint) for
+      renamed vars in the checkpoint
     :param TFNetwork network:
     """
     self.filename = filename
     self.network = network
     self.ignore_missing = ignore_missing
     self.params_prefix = params_prefix
     self.load_if_prefix = load_if_prefix
+    self.var_name_mapping = var_name_mapping or {}
     self.saveable_params = []
     count = 0
     for param in saveable_params:
@@ -3541,6 +3544,7 @@ def __init__(self, layer, checkpoint_loader):
       self.layer = layer
       self.prefix_param_name = layer.get_absolute_name_scope_prefix()
       self.checkpoint_param_names = []
+      self.var_name_mapping = checkpoint_loader.var_name_mapping
       prefix = self.prefix_param_name
       # Collect checkpoint params, and remove them from the lists.
       for name in list(checkpoint_loader.var_ckpt_names):
@@ -3576,7 +3580,7 @@ def assign_var(self, var, session):
         return
       self.assigned = True
       values_dict = {
-        name: self.reader.get_tensor(self.prefix_param_name + name)
+        name: self.reader.get_tensor(self.var_name_mapping.get(name, self.prefix_param_name + name))
         for name in self.checkpoint_param_names}
       self.reader = None  # Allow GC now, we do not need it anymore.
       print("Custom param import of layer %r with original params %r." % (
@@ -3894,6 +3898,7 @@ def get_lazy_dict(self):
       if v.endswith(MakeLoadCudnnRnn.cudnn_postfix):
         var_name_map.update(
           MakeLoadCudnnRnn(prefix=v[:-len(MakeLoadCudnnRnn.cudnn_postfix) + 1]).get_lazy_dict())
+    var_name_map.update({name: make_load_renamed(old_name) for name, old_name in self.var_name_mapping.items()})
 
     could_not_find_map_list = [v for v in missing_var_names if v not in var_name_map]
     if self.ignore_missing or not could_not_find_map_list: