Handle variables of unknown static shape (#738)

tensorpack/graph_builder/utils.py

@@ -79,6 +79,10 @@ class LeastLoadedDeviceSetter(object):
             self.ps_sizes), key=operator.itemgetter(1))
         device_name = self.ps_devices[device_index]
         var_size = op.outputs[0].get_shape().num_elements()
+        if var_size is None:
+            logger.warn("[LeastLoadedDeviceSetter] Shape of variable {} is not fully defined!".format(op.name))
+            var_size = 0
+
         self.ps_sizes[device_index] += var_size
 
         return sanitize_name(device_name)

tensorpack/tfutils/model_utils.py

@@ -28,9 +28,17 @@ def describe_trainable_vars():
             continue
         shape = v.get_shape()
         ele = shape.num_elements()
+        if ele is None:
+            logger.warn("Shape of variable {} is not fully defined but {}.".format(v.name, shape))
+            ele = 0
+        try:
+            shape = shape.as_list()
+        except ValueError:
+            shape = '<unknown>'
+
         total += ele
         total_bytes += ele * v.dtype.size
-        data.append([v.name, shape.as_list(), ele, v.device, v.dtype.base_dtype.name])
+        data.append([v.name, shape, ele, v.device, v.dtype.base_dtype.name])
     headers = ['name', 'shape', 'dim', 'device', 'dtype']
 
     dtypes = set([x[4] for x in data])