Skip to content

Commit 2cddc38

Browse files
mathpluscodemathpluscode
authored
Merge pull request #732 from DeepRegNet/revert-719-690-nan-inf-loss
Revert "690 nan inf loss"
2 parents e47c569 + adb3b95 commit 2cddc38

File tree

11 files changed

+67
-73
lines changed

11 files changed

+67
-73
lines changed

‎CHANGELOG.md‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ compatible with the updates.
2424

2525
### Changed
2626

27-
- Increased all EPS to 1e-5.
2827
- Clarify the suggestion in doc to use all-zero masks for missing labels.
2928
- Moved contributor list to a separate page.
3029
- Changed `no-test` flag to `full` for demo scripts.
@@ -38,7 +37,6 @@ compatible with the updates.
3837

3938
### Fixed
4039

41-
- Fixed LNCC loss regarding INF values.
4240
- Removed loss weight checks to be more robust.
4341
- Fixed import error under python 3.6.
4442
- Fixed the residual module in local net architecture, compatible for previous

‎deepreg/constant.py‎

Lines changed: 0 additions & 3 deletions
This file was deleted.

‎deepreg/dataset/loader/interface.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -379,10 +379,10 @@ def validate_images_and_labels(
379379
for arr, name in zip(
380380
[moving_image, fixed_image], ["moving_image", "fixed_image"]
381381
):
382-
if len(arr.shape) != 3 or min(arr.shape) <= 0:
382+
if len(arr.shape) != 3:
383383
raise ValueError(
384-
f"Sample {image_indices}'s {name}' shape should be 3D"
385-
f" and non-empty, got {arr.shape}."
384+
f"Sample {image_indices}'s {name}' shape should be 3D. "
385+
f"Got {arr.shape}."
386386
)
387387
# when data are labeled
388388
if moving_label is not None and fixed_label is not None:

‎deepreg/loss/image.py‎

Lines changed: 41 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
"""Provide different loss or metrics classes for images."""
22
import tensorflow as tf
33

4-
from deepreg.constant import EPS
54
from deepreg.loss.util import NegativeLossMixin
65
from deepreg.loss.util import gaussian_kernel1d_size as gaussian_kernel1d
76
from deepreg.loss.util import (
@@ -11,6 +10,8 @@
1110
)
1211
from deepreg.registry import REGISTRY
1312

13+
EPS = tf.keras.backend.epsilon()
14+
1415

1516
@REGISTRY.register_loss(name="ssd")
1617
class SumSquaredDifference(tf.keras.losses.Loss):
@@ -155,33 +156,27 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
155156
156157
E[t] = sum_i(w_i * t_i) / sum_i(w_i)
157158
158-
Here, we assume sum_i(w_i) == 1, means the weights have been normalized.
159-
160159
Similarly, the discrete variance in the window V[t] is
161160
162-
V[t] = E[(t - E[t])**2]
161+
V[t] = E[t**2] - E[t] ** 2
163162
164163
The local squared zero-normalized cross-correlation is therefore
165164
166165
E[ (t-E[t]) * (p-E[p]) ] ** 2 / V[t] / V[p]
167166
168-
When calculating variance, we choose to subtract the mean first then calculte
169-
variance instead of calculating E[t**2] - E[t] ** 2, the reason is that when
170-
E[t**2] and E[t] ** 2 are both very large or very small, the subtraction may
171-
have large rounding error and makes the result inaccurate. Also, it is not
172-
guaranteed that the result >= 0. For more details, please read "Algorithms for
173-
computing the sample variance: Analysis and recommendations." page 1.
167+
where the expectation in numerator is
168+
169+
E[ (t-E[t]) * (p-E[p]) ] = E[t * p] - E[t] * E[p]
170+
171+
Different kernel corresponds to different weights.
174172
175173
For now, y_true and y_pred have to be at least 4d tensor, including batch axis.
176174
177175
Reference:
178176
179177
- Zero-normalized cross-correlation (ZNCC):
180178
https://en.wikipedia.org/wiki/Cross-correlation
181-
- https://en.wikipedia.org/wiki/Weighted_arithmetic_mean#Reliability_weights
182-
- Chan, Tony F., Gene H. Golub, and Randall J. LeVeque.
183-
"Algorithms for computing the sample variance: Analysis and recommendations."
184-
The American Statistician 37.3 (1983): 242-247.
179+
- Code: https://github.com/voxelmorph/voxelmorph/blob/legacy/src/losses.py
185180
"""
186181

187182
kernel_fn_dict = dict(
@@ -217,8 +212,13 @@ def __init__(
217212
self.kernel_size = kernel_size
218213

219214
# (kernel_size, )
220-
# sum of the kernel weights would be one
221215
self.kernel = self.kernel_fn(kernel_size=self.kernel_size)
216+
# E[1] = sum_i(w_i), ()
217+
self.kernel_vol = tf.reduce_sum(
218+
self.kernel[:, None, None]
219+
* self.kernel[None, :, None]
220+
* self.kernel[None, None, :]
221+
)
222222

223223
def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
224224
"""
@@ -230,29 +230,38 @@ def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
230230
or (batch, dim1, dim2, dim3, ch)
231231
:return: shape = (batch,)
232232
"""
233-
# adjust shape to be (batch, dim1, dim2, dim3, ch)
233+
# adjust
234234
if len(y_true.shape) == 4:
235235
y_true = tf.expand_dims(y_true, axis=4)
236236
y_pred = tf.expand_dims(y_pred, axis=4)
237237
assert len(y_true.shape) == len(y_pred.shape) == 5
238238

239239
# t = y_true, p = y_pred
240-
t_mean = separable_filter(y_true, kernel=self.kernel)
241-
p_mean = separable_filter(y_pred, kernel=self.kernel)
242-
243-
t = y_true - t_mean
244-
p = y_pred - p_mean
245-
246-
# the variance can be biased but as both num and denom are biased
247-
# it got cancelled
248-
# https://en.wikipedia.org/wiki/Weighted_arithmetic_mean#Reliability_weights
249-
cross = separable_filter(t * p, kernel=self.kernel)
250-
t_var = separable_filter(t * t, kernel=self.kernel)
251-
p_var = separable_filter(p * p, kernel=self.kernel)
252-
253-
num = cross * cross
254-
denom = t_var * p_var
255-
ncc = (num + EPS) / (denom + EPS)
240+
# (batch, dim1, dim2, dim3, ch)
241+
t2 = y_true * y_true
242+
p2 = y_pred * y_pred
243+
tp = y_true * y_pred
244+
245+
# sum over kernel
246+
# (batch, dim1, dim2, dim3, 1)
247+
t_sum = separable_filter(y_true, kernel=self.kernel) # E[t] * E[1]
248+
p_sum = separable_filter(y_pred, kernel=self.kernel) # E[p] * E[1]
249+
t2_sum = separable_filter(t2, kernel=self.kernel) # E[tt] * E[1]
250+
p2_sum = separable_filter(p2, kernel=self.kernel) # E[pp] * E[1]
251+
tp_sum = separable_filter(tp, kernel=self.kernel) # E[tp] * E[1]
252+
253+
# average over kernel
254+
# (batch, dim1, dim2, dim3, 1)
255+
t_avg = t_sum / self.kernel_vol # E[t]
256+
p_avg = p_sum / self.kernel_vol # E[p]
257+
258+
# shape = (batch, dim1, dim2, dim3, 1)
259+
cross = tp_sum - p_avg * t_sum # E[tp] * E[1] - E[p] * E[t] * E[1]
260+
t_var = t2_sum - t_avg * t_sum # V[t] * E[1]
261+
p_var = p2_sum - p_avg * p_sum # V[p] * E[1]
262+
263+
# (E[tp] - E[p] * E[t]) ** 2 / V[t] / V[p]
264+
ncc = (cross * cross + EPS) / (t_var * p_var + EPS)
256265

257266
return tf.reduce_mean(ncc, axis=[1, 2, 3, 4])
258267

‎deepreg/loss/label.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,13 @@
44

55
import tensorflow as tf
66

7-
from deepreg.constant import EPS
87
from deepreg.loss.util import NegativeLossMixin, cauchy_kernel1d
98
from deepreg.loss.util import gaussian_kernel1d_sigma as gaussian_kernel1d
109
from deepreg.loss.util import separable_filter
1110
from deepreg.registry import REGISTRY
1211

12+
EPS = tf.keras.backend.epsilon()
13+
1314

1415
class MultiScaleLoss(tf.keras.losses.Loss):
1516
"""

‎deepreg/loss/util.py‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,25 +27,25 @@ def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
2727
return -super().call(y_true=y_true, y_pred=y_pred)
2828

2929

30+
EPS = tf.keras.backend.epsilon()
31+
32+
3033
def rectangular_kernel1d(kernel_size: int) -> tf.Tensor:
3134
"""
32-
Return a the 1D rectangular kernel for LocalNormalizedCrossCorrelation.
33-
34-
Sum of the weights is 1.
35+
Return a the 1D filter for separable convolution equivalent to a 3-D rectangular
36+
kernel for LocalNormalizedCrossCorrelation.
3537
3638
:param kernel_size: scalar, size of the 1-D kernel
3739
:return: kernel_weights, of shape (kernel_size, )
3840
"""
3941

40-
kernel = tf.ones(shape=(kernel_size,), dtype=tf.float32) / float(kernel_size)
42+
kernel = tf.ones(shape=(kernel_size,), dtype=tf.float32)
4143
return kernel
4244

4345

4446
def triangular_kernel1d(kernel_size: int) -> tf.Tensor:
4547
"""
46-
Return a the 1D triangular kernel for LocalNormalizedCrossCorrelation.
47-
48-
Sum of the weights is 1.
48+
1D triangular kernel.
4949
5050
Assume kernel_size is odd, it will be a smoothed from
5151
a kernel which center part is zero.
@@ -73,17 +73,13 @@ def triangular_kernel1d(kernel_size: int) -> tf.Tensor:
7373
kernel = tf.nn.conv1d(
7474
kernel[None, :, None], filters=filters, stride=[1, 1, 1], padding="SAME"
7575
)
76-
kernel = kernel / tf.reduce_sum(kernel)
77-
7876
return kernel[0, :, 0]
7977

8078

8179
def gaussian_kernel1d_size(kernel_size: int) -> tf.Tensor:
8280
"""
83-
Return a the 1D Gaussian kernel for LocalNormalizedCrossCorrelation.
84-
85-
Sum of the weights is 1.
86-
81+
Return a the 1D filter for separable convolution equivalent to a 3-D Gaussian
82+
kernel for LocalNormalizedCrossCorrelation.
8783
:param kernel_size: scalar, size of the 1-D kernel
8884
:return: filters, of shape (kernel_size, )
8985
"""
@@ -92,7 +88,6 @@ def gaussian_kernel1d_size(kernel_size: int) -> tf.Tensor:
9288

9389
grid = tf.range(0, kernel_size, dtype=tf.float32)
9490
filters = tf.exp(-tf.square(grid - mean) / (2 * sigma ** 2))
95-
filters = filters / tf.reduce_sum(filters)
9691

9792
return filters
9893

‎deepreg/model/network.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@ class RegistrationModel(tf.keras.Model):
3030

3131
def __init__(
3232
self,
33-
moving_image_size: Tuple,
34-
fixed_image_size: Tuple,
33+
moving_image_size: tuple,
34+
fixed_image_size: tuple,
3535
index_size: int,
3636
labeled: bool,
3737
batch_size: int,
@@ -61,7 +61,6 @@ def __init__(
6161
self.config = config
6262
self.num_devices = num_devices
6363
self.global_batch_size = num_devices * batch_size
64-
assert self.global_batch_size > 0
6564

6665
self._inputs = None # save inputs of self._model as dict
6766
self._outputs = None # save outputs of self._model as dict
@@ -223,6 +222,7 @@ def _build_loss(self, name: str, inputs_dict: dict):
223222

224223
# add loss
225224
self._model.add_loss(weighted_loss)
225+
226226
# add metric
227227
self._model.add_metric(
228228
loss_value, name=f"loss/{name}_{loss_layer.name}", aggregation="mean"

‎test/unit/test_interface.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,7 @@ def mock_sample_index_generator():
356356
fixed_label=None,
357357
image_indices=[1],
358358
)
359-
assert "Sample [1]'s moving_image' shape should be 3D" in str(err_info.value)
359+
assert "Sample [1]'s moving_image' shape should be 3D. " in str(err_info.value)
360360
with pytest.raises(ValueError) as err_info:
361361
generator.validate_images_and_labels(
362362
fixed_image=dummy_array,

‎test/unit/test_loss_label.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import tensorflow as tf
1313

1414
import deepreg.loss.label as label
15-
from deepreg.constant import EPS
1615

1716

1817
class TestMultiScaleLoss:
@@ -92,11 +91,11 @@ def y_pred(self):
9291
@pytest.mark.parametrize(
9392
"binary,background_weight,scales,expected",
9493
[
95-
(True, 0.0, None, -np.log(EPS)),
96-
(False, 0.0, None, -0.6 * np.log(0.3 + EPS)),
97-
(False, 0.2, None, -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
98-
(False, 0.2, [0, 0], -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
99-
(False, 0.2, [0, 1], 0.5239465),
94+
(True, 0.0, None, -np.log(1.0e-7)),
95+
(False, 0.0, None, -0.6 * np.log(0.3)),
96+
(False, 0.2, None, -0.48 * np.log(0.3) - 0.08 * np.log(0.7)),
97+
(False, 0.2, [0, 0], -0.48 * np.log(0.3) - 0.08 * np.log(0.7)),
98+
(False, 0.2, [0, 1], 0.5239637),
10099
],
101100
)
102101
def test_call(self, y_true, y_pred, binary, background_weight, scales, expected):
@@ -136,7 +135,7 @@ def y_pred(self):
136135
(True, None, 0),
137136
(False, None, 0.25),
138137
(False, [0, 0], 0.25),
139-
(False, [0, 1], 0.17485845),
138+
(False, [0, 1], 0.17484076),
140139
],
141140
)
142141
def test_call(self, y_true, y_pred, binary, scales, expected):

‎test/unit/test_loss_util.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,6 @@ def test_gaussian_kernel1d_size(kernel_size):
6262

6363
grid = tf.range(0, kernel_size, dtype=tf.float32)
6464
expected = tf.exp(-tf.square(grid - mean) / (2 * sigma ** 2))
65-
expected = expected / tf.reduce_sum(expected)
6665

6766
got = gaussian_kernel1d_size(kernel_size)
6867
assert is_equal_tf(got, expected)
@@ -76,7 +75,6 @@ def test_rectangular_kernel1d(kernel_size):
7675
:return:
7776
"""
7877
expected = tf.ones(shape=(kernel_size,), dtype=tf.float32)
79-
expected = expected / tf.reduce_sum(expected)
8078
got = rectangular_kernel1d(kernel_size)
8179
assert is_equal_tf(got, expected)
8280

@@ -93,7 +91,6 @@ def test_triangular_kernel1d(kernel_size):
9391
for it_k in range(kernel_size // 2):
9492
expected[it_k] = it_k + 1
9593
expected[-it_k - 1] = it_k + 1
96-
expected = expected / tf.reduce_sum(expected)
9794

9895
got = triangular_kernel1d(kernel_size)
9996
assert is_equal_tf(got, expected)

0 commit comments

Comments
 (0)