def eval_c_grad(layer, content):
nonlocal loss
feat = content.features[layer][:, start_[0]:end[0],
start_[1]:end[1]]
c_grad = self.data[layer] - feat
loss += lw * content_weight[layer] * norm2(c_grad)
axpy(lw * content_weight[layer], normalize(c_grad),
self.diff[layer])
def eval_s_grad(layer, style):
nonlocal loss
current_gram = gram_matrix(self.data[layer])
n, mh, mw = self.data[layer].shape
feat = self.data[layer].reshape((n, mh * mw))
gram_diff = current_gram - style.grams[layer]
s_grad = self._arr_pool.array_like(feat)
np.dot(gram_diff, feat, s_grad)
s_grad = s_grad.reshape((n, mh, mw))
loss += lw * style_weight[layer] * norm2(gram_diff) / len(
self.styles) / 2
axpy(lw * style_weight[layer] / len(self.styles),
normalize(s_grad), self.diff[layer])
def eval_loss_and_grad(self, img, sc_grad_args):
"""Returns the summed loss and gradient."""
old_img = self.model.img
self.model.img = img
lw = self.layer_weights['data']
# Compute style+content gradient
loss, grad = self.model.eval_sc_grad(*sc_grad_args)
# Compute total variation gradient
if ARGS.tv_weight:
tv_loss, tv_grad = tv_norm(self.model.img / 127.5,
beta=ARGS.tv_power)
loss += lw * ARGS.tv_weight * tv_loss
axpy(lw * ARGS.tv_weight, tv_grad, grad)
# Compute SWT norm and gradient
if ARGS.swt_weight:
swt_loss, swt_grad = swt_norm(self.model.img / 127.5,
ARGS.swt_wavelet,
ARGS.swt_levels,
p=ARGS.swt_power)
loss += lw * ARGS.swt_weight * swt_loss
axpy(lw * ARGS.swt_weight, swt_grad, grad)
# Compute p-norm regularizer gradient (from jcjohnson/cnn-vis and [3])
if ARGS.p_weight:
p_loss, p_grad = p_norm(
(self.model.img + self.model.mean - 127.5) / 127.5,
p=ARGS.p_power)
loss += lw * ARGS.p_weight * p_loss
axpy(lw * ARGS.p_weight, p_grad, grad)
# Compute auxiliary image gradient
if self.aux_image is not None:
aux_grad = (self.model.img - self.aux_image) / 127.5
loss += lw * ARGS.aux_weight * norm2(aux_grad)
axpy(lw * ARGS.aux_weight, aux_grad, grad)
self.model.img = old_img
return loss, grad
def eval_sc_grad_tile(self, img, start, layers, content_layers,
style_layers, dd_layers, layer_weights,
content_weight, style_weight, dd_weight):
"""Evaluates an individual style+content gradient tile."""
self.net.blobs['data'].reshape(1, 3, *img.shape[-2:])
self.data['data'] = img
loss = 0
# Prepare gradient buffers and run the model forward
for layer in layers:
self.diff[layer] = 0
self.net.forward(end=layers[0])
self.data[layers[0]] = np.maximum(0, self.data[layers[0]])
for i, layer in enumerate(layers):
lw = layer_weights[layer]
scale, _ = self.layer_info(layer)
start_ = start // scale
end = start_ + np.array(self.data[layer].shape[-2:])
def eval_c_grad(layer, content):
nonlocal loss
feat = content.features[layer][:, start_[0]:end[0],
start_[1]:end[1]]
c_grad = self.data[layer] - feat
loss += lw * content_weight[layer] * norm2(c_grad)
axpy(lw * content_weight[layer], normalize(c_grad),
self.diff[layer])
def eval_s_grad(layer, style):
nonlocal loss
current_gram = gram_matrix(self.data[layer])
n, mh, mw = self.data[layer].shape
feat = self.data[layer].reshape((n, mh * mw))
gram_diff = current_gram - style.grams[layer]
s_grad = self._arr_pool.array_like(feat)
np.dot(gram_diff, feat, s_grad)
s_grad = s_grad.reshape((n, mh, mw))
loss += lw * style_weight[layer] * norm2(gram_diff) / len(
self.styles) / 2
axpy(lw * style_weight[layer] / len(self.styles),
normalize(s_grad), self.diff[layer])
# Compute the content and style gradients
if layer in content_layers:
for content in self.contents:
eval_c_grad(layer, content)
if layer in style_layers:
for style in self.styles:
eval_s_grad(layer, style)
if layer in dd_layers:
loss -= lw * dd_weight[layer] * norm2(self.data[layer])
axpy(-lw * dd_weight[layer], normalize(self.data[layer]),
self.diff[layer])
# Run the model backward
if i + 1 == len(layers):
self.net.backward(start=layer)
else:
self.net.backward(start=layer, end=layers[i + 1])
return loss, self.diff['data']