Nan/inf loss encountered

[mathpluscode]

Subject of the issue

More discussion is here Project-MONAI/MONAI#1868

Using script https://github.com/DeepRegNet/Benchmark/blob/9-data-for-miccai/benchmark/balakrishnan2019/voxel_morph_balakrishnan_2019.py with the latest main branch.

226/226 [==============================] - 3083s 14s/step - loss: -0.7058 - loss/regularization_GradientNorm: 2.1987e-04 - loss/regularization_GradientNorm_weighted: 2.1987e-04 - loss/image_LocalNormalizedCrossCorrelationLoss: -0.7060 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.7060
Epoch 2/2000
226/226 [==============================] - 3266s 14s/step - loss: -0.4778 - loss/regularization_GradientNorm: 0.3009 - loss/regularization_GradientNorm_weighted: 0.3009 - loss/image_LocalNormalizedCrossCorrelationLoss: -0.7787 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.7787
Epoch 3/2000
226/226 [==============================] - 3280s 15s/step - loss: nan - loss/regularization_GradientNorm: nan - loss/regularization_GradientNorm_weighted: nan - loss/image_LocalNormalizedCrossCorrelationLoss: -inf - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -inf

Your environment

• DeepReg version (commit hash)

  5e5bbdc

• OS (e.g. Ubuntu 20.04, MacOS 10.15, etc.)

  Ubuntu

• Python Version (3.7, 3.8, etc.)

  3.7

• TensorFlow

  2.3

[mathpluscode]

The loss is not nan/inf when using the following config

dataset:
  dir:
    train: "/raid/candi/Yunguan/DeepReg/neuroimaging/preprocessed" # required
    valid:
    test:
  format: "nifti"
  type: "unpaired" # paired / unpaired / grouped
  labeled: false # whether to use the labels if available, "true" or "false"
  image_shape: [196, 112, 96]

train:
  # define neural network structure
  method: "ddf" # options include "ddf", "dvf", "conditional"
  backbone:
    name: "unet" # options include "local", "unet" and "global"
    num_channel_initial: 16 # number of initial channel in local net, controls the size of the network
    depth: 4
    # concat_skip: true
    # encode_num_channels: [16, 32, 32, 32, 32]
    # decode_num_channels: [32, 32, 32, 32, 32]

  # define the loss function for training
  loss:
    image:
      name: "lncc" # other options include "lncc", "ssd" and "gmi", for local normalised cross correlation,
      weight: 1.0
    label:
      weight: 0.0
      name: "dice" # options include "dice", "cross-entropy", "mean-squared", "generalised_dice" and "jaccard"
    regularization:
      weight: 1.0 # weight of regularization loss
      name: "gradient" # options include "bending", "gradient"

  # define the optimizer
  optimizer:
    name: "adam" # options include "adam", "sgd" and "rms"
    adam:
      learning_rate: 1.0e-4

  # define the hyper-parameters for preprocessing
  preprocess:
    data_augmentation:
      name: "affine"
    batch_size: 4
    shuffle_buffer_num_batch: 1 # shuffle_buffer_size = batch_size * shuffle_buffer_num_batch

  # other training hyper-parameters
  epochs: 2000 # number of training epochs
  save_period: 50 # the model will be saved every `save_period` epochs.

[mathpluscode]

226/226 [==============================] - 4779s 21s/step - loss: -0.1736 - loss/regularization_GradientNorm: 4.3662e-05 - loss/regularization_GradientNorm_weighted: 4.3662e-05 - loss/image_LocalNormalizedCrossC
orrelationLoss: -0.1736 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.1736
Epoch 2/50
226/226 [==============================] - 5154s 23s/step - loss: -0.1405 - loss/regularization_GradientNorm: 0.0362 - loss/regularization_GradientNorm_weighted: 0.0362 - loss/image_LocalNormalizedCrossCorrelati
onLoss: -0.1767 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.1767
Epoch 3/50
226/226 [==============================] - 5218s 23s/step - loss: -0.1474 - loss/regularization_GradientNorm: 0.0315 - loss/regularization_GradientNorm_weighted: 0.0315 - loss/image_LocalNormalizedCrossCorrelati
onLoss: -0.1788 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.1788
Epoch 4/50
226/226 [==============================] - 5211s 23s/step - loss: nan - loss/regularization_GradientNorm: nan - loss/regularization_GradientNorm_weighted: nan - loss/image_LocalNormalizedCrossCorrelationLoss: -i
nf - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -inf
Epoch 5/50
226/226 [==============================] - ETA: 0s - loss: nan - loss/regularization_GradientNorm: nan - loss/regularization_GradientNorm_weighted: nan - loss/image_LocalNormalizedCrossCorrelationLoss: -0.2500 -
 loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.2500 2021-03-15 06:06:09.605122: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at summary_kernels.cc:242 : Invalid argument: Nan i
n summary histogram for: Unet/conv3d/kernel_0
Traceback (most recent call last):
  File "voxel_morph_balakrishnan_2019.py", line 217, in <module>
    main()
  File "voxel_morph_balakrishnan_2019.py", line 212, in main
    ckpt_path="",
  File "/raid/candi/yunguan/Git/DeepReg/deepreg/train.py", line 167, in train
    callbacks=callbacks,
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 108, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 1137, in fit
    callbacks.on_epoch_end(epoch, epoch_logs)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/keras/callbacks.py", line 412, in on_epoch_end
    callback.on_epoch_end(epoch, logs)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/keras/callbacks.py", line 2182, in on_epoch_end
    self._log_weights(epoch)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/keras/callbacks.py", line 2234, in _log_weights
    summary_ops_v2.histogram(weight_name, weight, step=epoch)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/ops/summary_ops_v2.py", line 836, in histogram
    return summary_writer_function(name, tensor, function, family=family)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/ops/summary_ops_v2.py", line 765, in summary_writer_function
    should_record_summaries(), record, _nothing, name="")
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/framework/smart_cond.py", line 54, in smart_cond
    return true_fn()
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/ops/summary_ops_v2.py", line 758, in record
    with ops.control_dependencies([function(tag, scope)]):
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/ops/summary_ops_v2.py", line 834, in function
    name=scope)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/ops/gen_summary_ops.py", line 480, in write_histogram_summary
    writer, step, tag, values, name=name, ctx=_ctx)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/ops/gen_summary_ops.py", line 499, in write_histogram_summary_eager_fallback
    attrs=_attrs, ctx=ctx, name=name)
  File "/home/yunguan/miniconda3/envs/deepreg/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Nan in summary histogram for: Unet/conv3d/kernel_0 [Op:WriteHistogramSummary]

[mathpluscode]

It seems that the problem may come from the deformation regularization.

[mathpluscode]

changing deform loss to https://github.com/voxelmorph/voxelmorph/blob/legacy/src/losses.py#L88 does not solve the problem

Need to double check the resampler and related functions

eventually we got inf first in lncc, but the nan/inf checks in LNCC was not triggered?

[mathpluscode]

changed zero_boundary in resample, default to False, still got nan/inf

maybed linked to leaky relu

[mathpluscode]

change from leaky relu to relu does not solve the problem neither

[mathpluscode]

ok just found out that I added activation at the final layer before outputting DDF. :( it might be the cause, as leaky relu or relu both push DDF towards positive values, which may push DDF to be too large

[mathpluscode]

by fixing activation, we still got nan/inf

[mathpluscode]

change to SSD loss for image, then the nan/inf is gone.

So the problem must come from the LNCC loss