-
Notifications
You must be signed in to change notification settings - Fork 0
/
engine.py
185 lines (149 loc) · 5.66 KB
/
engine.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
r"""Engine for iterative optimization-based registration."""
# TODO: Use ignite.engine instead of custom RegistrationEngine.
from __future__ import annotations
import math
import pickle
from collections import OrderedDict
from timeit import default_timer as timer
from typing import Callable
import matplotlib.pyplot as plt
import torch
from torch import Tensor
from torch.optim import Optimizer
from torch.utils.hooks import RemovableHandle
from optim import slope_of_least_squares_fit
from losses import RegistrationLoss, RegistrationResult
PROFILING = False
class RegistrationEngine(object):
r"""Minimize registration loss until convergence."""
def __init__(
self,
loss: RegistrationLoss,
optimizer: Optimizer,
max_steps: int = 500,
min_delta: float = 1e-6,
min_value: float = float("nan"),
max_history: int = 1000,
):
r"""Initialize registration loop."""
self.loss = loss
self.optimizer = optimizer
self.num_steps = 0
self.max_steps = max_steps
self.min_delta = min_delta
self.min_value = min_value
self.max_history = max(2, max_history)
self.loss_values = []
self._eval_hooks = OrderedDict()
self._step_hooks = OrderedDict()
@property
def loss_value(self) -> float:
if not self.loss_values:
return float("inf")
return self.loss_values[-1]
def step(self) -> float:
r"""Perform one registration step.
Returns:
Loss value prior to taking gradient step.
"""
num_evals = 0
def closure() -> float:
self.optimizer.zero_grad()
t_start = timer()
result = self.loss.eval()
if PROFILING:
print(f"Forward pass in {timer() - t_start:.3f}s")
loss = result["loss"]
assert isinstance(loss, Tensor)
t_start = timer()
loss.backward()
if PROFILING:
print(f"Backward pass in {timer() - t_start:.3f}s")
nonlocal num_evals
num_evals += 1
with torch.no_grad():
for hook in self._eval_hooks.values():
hook(self, self.num_steps, num_evals, result)
return float(loss)
loss_value = self.optimizer.step(closure)
assert loss_value is not None
with torch.no_grad():
for hook in self._step_hooks.values():
hook(self, self.num_steps, num_evals, loss_value)
return loss_value
def run(self) -> float:
r"""Perform registration steps until convergence.
Returns:
Loss value prior to taking last gradient step.
"""
self.loss_values = []
self.num_steps = 0
while self.num_steps < self.max_steps and not self.converged():
value = self.step()
self.num_steps += 1
if math.isnan(value):
raise RuntimeError(
f"NaN value in registration loss at gradient step {self.num_steps}"
)
if math.isinf(value):
raise RuntimeError(
f"Inf value in registration loss at gradient step {self.num_steps}"
)
self.loss_values.append(value)
if len(self.loss_values) > self.max_history:
self.loss_values.pop(0)
return self.loss_value
def converged(self) -> bool:
r"""Check convergence criteria."""
values = self.loss_values
if not values:
return False
value = values[-1]
if self.min_delta < 0:
epsilon = abs(self.min_delta * value)
else:
epsilon = self.min_delta
slope = slope_of_least_squares_fit(values)
if abs(slope) < epsilon:
return True
if value < self.min_value:
return True
return False
def register_eval_hook(
self, hook: Callable[[RegistrationEngine, int, int, RegistrationResult], None]
) -> RemovableHandle:
r"""Registers a evaluation hook.
The hook will be called every time after the registration loss has been evaluated
during a single step of the optimizer, and the backward pass was performed, but
before the model parameters are updated by taking a gradient step.
hook(self, num_steps: int, num_evals: int, result: RegistrationResult) -> None
Returns:
A handle that can be used to remove the added hook by calling ``handle.remove()``
"""
handle = RemovableHandle(self._eval_hooks)
self._eval_hooks[handle.id] = hook
return handle
def register_step_hook(
self, hook: Callable[[RegistrationEngine, int, int, float], None]
) -> RemovableHandle:
r"""Registers a gradient step hook.
The hook will be called every time after a gradient step of the optimizer.
hook(self, num_steps: int, num_evals: int, loss: float) -> None
Returns:
A handle that can be used to remove the added hook by calling ``handle.remove()``
"""
handle = RemovableHandle(self._step_hooks)
self._step_hooks[handle.id] = hook
return handle
def plot_loss(self, outfile):
plt.figure(figsize=(10, 6))
plt.plot(range(len(self.loss_values)), self.loss_values, marker='o')
plt.title('Loss Values Over Steps')
plt.xlabel('Steps')
plt.ylabel('Loss Value')
plt.grid(True)
plt.savefig(outfile)
plt.show()
def dump_loss_values(self, outfile):
with open(outfile, 'wb') as f:
pickle.dump(self.loss_values, f)