Fix mirrored variable load (#895)

* fix load variable bug when use mirrored strategy

* add custom aspect ratio support

* fix bug

* optimize performance

Author: fsx950223

efficientdet/keras/anchors.py

@@ -116,8 +116,11 @@ def _generate_boxes(self):
         stride, octave_scale, aspect, anchor_scale = config
         base_anchor_size_x = anchor_scale * stride[1] * 2**octave_scale
         base_anchor_size_y = anchor_scale * stride[0] * 2**octave_scale
-        aspect_x = np.sqrt(aspect)
-        aspect_y = 1.0 / aspect_x
+        if isinstance(aspect, list):
+          aspect_x, aspect_y = aspect
+        else:
+          aspect_x = np.sqrt(aspect)
+          aspect_y = 1.0 / aspect_x
         anchor_size_x_2 = base_anchor_size_x * aspect_x / 2.0
         anchor_size_y_2 = base_anchor_size_y * aspect_y / 2.0
 

efficientdet/utils.py

@@ -111,35 +111,47 @@ def get_ckpt_var_map(ckpt_path, ckpt_scope, var_scope, skip_mismatch=None):
   ckpt_var_name_to_shape = reader.get_variable_to_shape_map()
   ckpt_var_names = set(reader.get_variable_to_shape_map().keys())
 
+  if tf.distribute.get_replica_context():
+    replica_id = tf.keras.backend.get_value(
+      tf.distribute.get_replica_context().replica_id_in_sync_group)
+  else:
+    replica_id = 0
+
   for i, v in enumerate(model_vars):
-    if not v.op.name.startswith(var_scope):
+    var_op_name = v.op.name
+
+    if replica_id >= 1:
+      var_op_name = ''.join(var_op_name.rsplit(f'/replica_{replica_id}', 1))
+
+    if not var_op_name.startswith(var_scope):
       logging.info('skip {} -- does not match scope {}'.format(
-          v.op.name, var_scope))
-    ckpt_var = ckpt_scope + v.op.name[len(var_scope):]
+        var_op_name, var_scope))
+    ckpt_var = ckpt_scope + var_op_name[len(var_scope):]
+
     if (ckpt_var not in ckpt_var_names and
-        v.op.name.endswith('/ExponentialMovingAverage')):
-      ckpt_var = ckpt_scope + v.op.name[:-len('/ExponentialMovingAverage')]
+              var_op_name.endswith('/ExponentialMovingAverage')):
+      ckpt_var = ckpt_scope + var_op_name[:-len('/ExponentialMovingAverage')]
 
     if ckpt_var not in ckpt_var_names:
       if 'Momentum' in ckpt_var or 'RMSProp' in ckpt_var:
         # Skip optimizer variables.
         continue
       if skip_mismatch:
-        logging.info('skip {} ({}) -- not in ckpt'.format(v.op.name, ckpt_var))
+        logging.info('skip {} ({}) -- not in ckpt'.format(var_op_name, ckpt_var))
         continue
       raise ValueError('{} is not in ckpt {}'.format(v.op, ckpt_path))
 
     if v.shape != ckpt_var_name_to_shape[ckpt_var]:
       if skip_mismatch:
         logging.info('skip {} ({} vs {}) -- shape mismatch'.format(
-            v.op.name, v.shape, ckpt_var_name_to_shape[ckpt_var]))
+          var_op_name, v.shape, ckpt_var_name_to_shape[ckpt_var]))
         continue
       raise ValueError('shape mismatch {} ({} vs {})'.format(
-          v.op.name, v.shape, ckpt_var_name_to_shape[ckpt_var]))
+        var_op_name, v.shape, ckpt_var_name_to_shape[ckpt_var]))
 
     if i < 5:
       # Log the first few elements for sanity check.
-      logging.info('Init {} from ckpt var {}'.format(v.op.name, ckpt_var))
+      logging.info('Init {} from ckpt var {}'.format(var_op_name, ckpt_var))
     var_map[ckpt_var] = v
 
   return var_map