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_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']