Error loading trained models

Hello author(s), Amazing work and thank you for sharing your code! I am trying to extend your work by training models using Spinal CT images and I am running into a problem when loading the models trained with your code. I am using .nii files scaled [0, 1] and of size (128, 128, 96) to train and have seemingly successfully run 1500 epochs, however, only the model trained on the first epoch will load when running the register.py file. All others fail with the following error:

  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/ops/array_ops.py", line 1694, in ones
    tensor_shape.TensorShape(shape))
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 250, in _tensor_shape_tensor_conversion_function
    "Cannot convert a partially known TensorShape to a Tensor: %s" % s)
ValueError: Cannot convert a partially known TensorShape to a Tensor: (1, ?, ?, ?, 3)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "register.py", line 103, in <module>
    register(**vars(args))
  File "register.py", line 66, in register
    net = keras.models.load_model(model_file, custom_objects=custom_objects)
  File "/home/phoenixyuwilkie/.local/lib/python3.6/site-packages/keras/engine/saving.py", line 419, in load_model
    model = _deserialize_model(f, custom_objects, compile)
  File "/home/phoenixyuwilkie/.local/lib/python3.6/site-packages/keras/engine/saving.py", line 312, in _deserialize_model
    sample_weight_mode=sample_weight_mode)
  File "/home/phoenixyuwilkie/.local/lib/python3.6/site-packages/keras/engine/training.py", line 342, in compile
    sample_weight, mask)
  File "/home/phoenixyuwilkie/.local/lib/python3.6/site-packages/keras/engine/training_utils.py", line 404, in weighted
    score_array = fn(y_true, y_pred)
  File "/home/phoenixyuwilkie/Documents/Sam/obl_voxelmorph/src/losses.py", line 221, in kl_loss
    self.D = self._degree_matrix(self.flow_vol_shape)
  File "/home/phoenixyuwilkie/Documents/Sam/obl_voxelmorph/src/losses.py", line 168, in _degree_matrix
    z = K.ones([1] + sz)
  File "/home/phoenixyuwilkie/.local/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py", line 734, in ones
    v = tf.ones(shape=shape, dtype=tf_dtype, name=name)
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/ops/array_ops.py", line 1697, in ones
    shape = ops.convert_to_tensor(shape, dtype=dtypes.int32)
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1050, in convert_to_tensor
    as_ref=False)
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1146, in internal_convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 229, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 208, in constant
    value, dtype=dtype, shape=shape, verify_shape=verify_shape))
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/tensor_util.py", line 442, in make_tensor_proto
    _AssertCompatible(values, dtype)
  File "/home/phoenixyuwilkie/miniconda3/envs/obl_tf/lib/python3.6/site-packages/tensorflow/python/framework/tensor_util.py", line 353, in _AssertCompatible
    (dtype.name, repr(mismatch), type(mismatch).__name__))
TypeError: Expected int32, got None of type '_Message' instead.

The training seems to have completed, I have attached the output from running the train_miccai2018.py file. nohup.txt

Thank you, Samuel Newhook

[JesperDramsch]

I ran into a similar problem when using register.py, after some investigation, I found that the MICCAI version is trained with:

flow_vol_shape = model.outputs[-1].shape[1:-1]
loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape)

and the custom_layer defintion contains:

custom_objects = { [......]
             'recon_loss': losses.Miccai2018(0.02, 10).recon_loss, # values shouldn't matter
             'kl_loss': losses.Miccai2018(0.02, 10).kl_loss        # values shouldn't matter
             }

The values don’t matter, but the type of the values matters, it seems with the most recent version on Github. The following code fixed the error “TypeError: Expected int32, got None of type ‘_Message’ instead.”.

custom_objects = { [......]
             'recon_loss': losses.Miccai2018(0.02, 10, flow_vol_shape=[32]).recon_loss, # values shouldn't matter
             'kl_loss': losses.Miccai2018(0.02, 10, flow_vol_shape=[32]).kl_loss        # values shouldn't matter
             }

[John1231983]

I just use original code without modify anything in train.py

train.txt