forked from Sohl-Dickstein/Diffusion-Probabilistic-Models
-
Notifications
You must be signed in to change notification settings - Fork 0
/
extensions.py
183 lines (157 loc) · 7.49 KB
/
extensions.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
"""
Extensions called during training to generate samples and diagnostic plots and printouts.
"""
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import numpy as np
import os
import theano.tensor as T
import theano
from blocks.extensions import SimpleExtension
import viz
import sampler
class PlotSamples(SimpleExtension):
def __init__(self, model, algorithm, X, path, n_samples=49, **kwargs):
"""
Generate samples from the model. The do() function is called as an extension during training.
Generates 3 types of samples:
- Sample from generative model
- Sample from image denoising posterior distribution (default signal to noise of 1)
- Sample from image inpainting posterior distribution (inpaint left half of image)
"""
super(PlotSamples, self).__init__(**kwargs)
self.model = model
self.path = path
self.X = X[:n_samples].reshape(
(n_samples, model.n_colors, model.spatial_width, model.spatial_width))
self.n_samples = n_samples
X_noisy = T.tensor4('X noisy samp')
t = T.matrix('t samp')
self.get_mu_sigma = theano.function([X_noisy, t], model.get_mu_sigma(X_noisy, t),
allow_input_downcast=True)
def do(self, callback_name, *args):
print "generating samples"
base_fname_part1 = self.path + '/samples-'
base_fname_part2 = '_epoch%04d'%self.main_loop.status['epochs_done']
sampler.generate_samples(self.model, self.get_mu_sigma,
n_samples=self.n_samples, inpaint=False, denoise_sigma=None, X_true=None,
base_fname_part1=base_fname_part1, base_fname_part2=base_fname_part2)
sampler.generate_samples(self.model, self.get_mu_sigma,
n_samples=self.n_samples, inpaint=True, denoise_sigma=None, X_true=self.X,
base_fname_part1=base_fname_part1, base_fname_part2=base_fname_part2)
sampler.generate_samples(self.model, self.get_mu_sigma,
n_samples=self.n_samples, inpaint=False, denoise_sigma=1, X_true=self.X,
base_fname_part1=base_fname_part1, base_fname_part2=base_fname_part2)
class PlotParameters(SimpleExtension):
def __init__(self, model, blocks_model, path, **kwargs):
super(PlotParameters, self).__init__(**kwargs)
self.path = path
self.model = model
self.blocks_model = blocks_model
def do(self, callback_name, *args):
print "plotting parameters"
for param_name, param in self.blocks_model.params.iteritems():
filename_safe_name = '-'.join(param_name.split('/')[2:]).replace(' ', '_')
base_fname_part1 = self.path + '/params-' + filename_safe_name
base_fname_part2 = '_epoch%04d'%self.main_loop.status['epochs_done']
viz.plot_parameter(param.get_value(), base_fname_part1, base_fname_part2,
title=param_name, n_colors=self.model.n_colors)
class PlotGradients(SimpleExtension):
def __init__(self, model, blocks_model, algorithm, X, path, **kwargs):
super(PlotGradients, self).__init__(**kwargs)
self.path = path
self.X = X
self.model = model
self.blocks_model = blocks_model
gradients = []
for param_name in sorted(self.blocks_model.params.keys()):
gradients.append(algorithm.gradients[self.blocks_model.params[param_name]])
self.grad_f = theano.function(algorithm.inputs, gradients, allow_input_downcast=True)
def do(self, callback_name, *args):
print "plotting gradients"
grad_vals = self.grad_f(self.X)
keynames = sorted(self.blocks_model.params.keys())
for ii in xrange(len(keynames)):
param_name = keynames[ii]
val = grad_vals[ii]
filename_safe_name = '-'.join(param_name.split('/')[2:]).replace(' ', '_')
base_fname_part1 = self.path + '/grads-' + filename_safe_name
base_fname_part2 = '_epoch%04d'%self.main_loop.status['epochs_done']
viz.plot_parameter(val, base_fname_part1, base_fname_part2,
title="grad " + param_name, n_colors=self.model.n_colors)
class PlotInternalState(SimpleExtension):
def __init__(self, model, blocks_model, state, features, X, path, **kwargs):
super(PlotInternalState, self).__init__(**kwargs)
self.path = path
self.X = X
self.model = model
self.blocks_model = blocks_model
self.internal_state_f = theano.function([features], state, allow_input_downcast=True)
self.internal_state_names = []
for var in state:
self.internal_state_names.append(var.name)
def do(self, callback_name, *args):
print "plotting internal state of network"
state = self.internal_state_f(self.X)
for ii in xrange(len(state)):
param_name = self.internal_state_names[ii]
val = state[ii]
filename_safe_name = param_name.replace(' ', '_').replace('/', '-')
base_fname_part1 = self.path + '/state-' + filename_safe_name
base_fname_part2 = '_epoch%04d'%self.main_loop.status['epochs_done']
viz.plot_parameter(val, base_fname_part1, base_fname_part2,
title="state " + param_name, n_colors=self.model.n_colors)
class PlotMonitors(SimpleExtension):
def __init__(self, path, burn_in_iters=0, **kwargs):
super(PlotMonitors, self).__init__(**kwargs)
self.path = path
self.burn_in_iters = burn_in_iters
def do(self, callback_name, *args):
print "plotting monitors"
try:
df = self.main_loop.log.to_dataframe()
except AttributeError:
# This starting breaking after a Blocks update.
print "Failed to generate monitoring plots due to Blocks interface change."
return
iter_number = df.tail(1).index
# Throw out the first burn_in values
# as the objective is often much larger
# in that period.
if iter_number > self.burn_in_iters:
df = df.loc[self.burn_in_iters:]
cols = [col for col in df.columns if col.startswith(('cost', 'train', 'test'))]
df = df[cols].interpolate(method='linear')
# If we don't have any non-nan dataframes, don't plot
if len(df) == 0:
return
try:
axs = df.interpolate(method='linear').plot(
subplots=True, legend=False, figsize=(5, len(cols)*2))
except TypeError:
# This starting breaking after a different Blocks update.
print "Failed to generate monitoring plots due to Blocks interface change."
return
for ax, cname in zip(axs, cols):
ax.set_title(cname)
fn = os.path.join(self.path,
'monitors_subplots_epoch%04d.png' % self.main_loop.status['epochs_done'])
plt.savefig(fn, bbox_inches='tight')
plt.clf()
df.plot(subplots=False, figsize=(15,10))
plt.gcf().tight_layout()
fn = os.path.join(self.path,
'monitors_epoch%04d.png' % self.main_loop.status['epochs_done'])
plt.savefig(fn, bbox_inches='tight')
plt.close('all')
def decay_learning_rate(iteration, old_value):
# TODO the numbers in this function should not be hard coded
# this is called every epoch
# reduce the learning rate by 10 every 1000 epochs
decay_rate = np.exp(np.log(0.1)/1000.)
new_value = decay_rate*old_value
if new_value < 1e-5:
new_value = 1e-5
print "learning rate %g"%new_value
return np.float32(new_value)