def train_step(batch):
with tf.GradientTape() as tape:
output_batch = it_network(batch, training=True)
output_batch = 255 * (output_batch + 1.0) / 2.0 # float deprocess
# Feed target and output batch through loss_network
target_batch_feature_maps = loss_network(batch)
output_batch_feature_maps = loss_network(output_batch)
c_loss = content_loss(
target_batch_feature_maps[hparams['content_layer_index']],
output_batch_feature_maps[hparams['content_layer_index']])
c_loss *= hparams['content_weight']
# Get output gram_matrix
output_gram_matrices = [
gram_matrix(x) for x in output_batch_feature_maps
]
s_loss = style_loss(target_gram_matrices, output_gram_matrices)
s_loss *= hparams['style_weight'] / num_style_layers
total_loss = c_loss + s_loss
scaled_loss = optimizer.get_scaled_loss(total_loss)
scaled_gradients = tape.gradient(scaled_loss,
it_network.trainable_variables)
gradients = optimizer.get_unscaled_gradients(scaled_gradients)
#gradients = tape.gradient(total_loss, it_network.trainable_variables)
optimizer.apply_gradients(
zip(gradients, it_network.trainable_variables))
total_loss_avg(total_loss)
content_loss_avg(c_loss)
style_loss_avg(s_loss)
def train_step(content_img, style_img):
t = transformer.encode(content_img, style_img, alpha=1.0)
with tf.GradientTape() as tape:
stylized_img = transformer.decode(t)
_, style_feat_style = vgg(style_img)
content_feat_stylized, style_feat_stylized = vgg(stylized_img)
tot_style_loss = args.style_weight * style_loss(
style_feat_style, style_feat_stylized)
tot_content_loss = args.content_weight * content_loss(
t, content_feat_stylized)
loss = tot_style_loss + tot_content_loss
gradients = tape.gradient(loss, transformer.trainable_variables)
optimizer.apply_gradients(
zip(gradients, transformer.trainable_variables))
train_loss(loss)
train_style_loss(tot_style_loss)
train_content_loss(tot_content_loss)
def train_step(images):
with tf.GradientTape() as tape:
transformed_images = transformer(images)
_, content_features = extractor(images)
style_features_transformed, content_features_transformed = extractor(
transformed_images)
tot_style_loss = args.style_weight * style_loss(
gram_style, style_features_transformed)
tot_content_loss = args.content_weight * content_loss(
content_features, content_features_transformed)
loss = tot_style_loss + tot_content_loss
gradients = tape.gradient(loss, transformer.trainable_variables)
optimizer.apply_gradients(
zip(gradients, transformer.trainable_variables))
train_loss(loss)
train_style_loss(tot_style_loss)
train_content_loss(tot_content_loss)
layer_features = outputs_dict['block5_conv2']
content_image_features = layer_features[0, :, :, :]
generated_features = layer_features[2, :, :, :]
loss += content_weight * utils.content_loss(content_image_features,
generated_features)
feature_layers = [
'block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1',
'block5_conv1'
]
for layer_name in feature_layers:
layer_features = outputs_dict[layer_name]
style_features = layer_features[1, :, :, :]
generated_features = layer_features[2, :, :, :]
sl = utils.style_loss(style_features, generated_features, num_rows,
num_cols)
loss += (style_weight / len(feature_layers)) * sl
loss += total_weight * utils.total_loss(generated_image, num_rows, num_cols)
# Gradients of the generated image wrt the loss
grads = K.gradients(loss, generated_image)
outputs = [loss]
if isinstance(grads, (list, tuple)):
outputs += grads
else:
outputs.append(grads)
f_outputs = K.function([generated_image], outputs)
x = proc.preprocess_image(content_image_path, num_rows, num_cols)
matrix[0, 2] = -np.sin(angle)
matrix[1, 1] = np.sin(angle)
matrix[2, 1] = np.cos(angle)
matrix[2, 2] = 1.
# print(matrix)
return matrix
def pad_matrix(matrix):
return np.hstack([matrix, np.zeros((3, 1))])
samples_x = samples[:, 0]
samples_y = samples[:, 1]
samples_z = samples[:, 2]
for i in range(batch_size):
theta[i] = pad_matrix(
np.dot(np.dot(rot_z(samples_z[i]), rot_y(samples_y[i])),
rot_x(samples_x[i])))
theta = torch.from_numpy(theta).float()
z = torch.randn(batch_size, 128)
G = Generator()
img = G(z, theta)
print(img.size())
D = Discriminator()
rf, z, logits = D(img)
print(rf.size(), z.size(), logits[0].size(), logits[1].size(),
logits[2].size(), logits[3].size())
from utils import style_loss
print(style_loss(logits, logits))