def get_region_op(region, region_weight):
return ops.MultiLayerEncodingOperator(
multi_layer_encoder,
style_layers,
get_style_op,
score_weight=region_weight,
)
def _get_perceptual_loss(
self,
*,
backbone: str,
content_layer: str,
content_weight: float,
style_layers: Sequence[str],
style_weight: float,
) -> loss.PerceptualLoss:
mle, _ = cast(enc.MultiLayerEncoder, self.backbones.get(backbone)())
content_loss = ops.FeatureReconstructionOperator(
mle.extract_encoder(content_layer), score_weight=content_weight)
style_loss = ops.MultiLayerEncodingOperator(
mle,
style_layers,
lambda encoder, layer_weight: self._modified_gram_loss(
encoder, score_weight=layer_weight),
layer_weights="sum",
score_weight=style_weight,
)
return loss.PerceptualLoss(content_loss, style_loss)
def test_MultiLayerEncodingOperator(self):
class TestOperator(ops.EncodingRegularizationOperator):
def input_enc_to_repr(self, image):
pass
def calculate_score(self, input_repr):
pass
def get_encoding_op(encoder, score_weight):
return TestOperator(encoder, score_weight=score_weight)
layers = [str(index) for index in range(3)]
modules = [(layer, nn.Module()) for layer in layers]
multi_layer_encoder = MultiLayerEncoder(modules)
multi_layer_enc_op = ops.MultiLayerEncodingOperator(
multi_layer_encoder, layers, get_encoding_op)
for layer in layers:
op = getattr(multi_layer_enc_op, layer)
self.assertIsInstance(op.encoder, SingleLayerEncoder)
self.assertEqual(op.encoder.layer, layer)
self.assertIs(op.encoder.multi_layer_encoder, multi_layer_encoder)
content_encoder = multi_layer_encoder.extract_encoder(content_layer)
content_weight = 1e0
content_loss = ops.FeatureReconstructionOperator(content_encoder,
score_weight=content_weight)
style_layers = ("relu1_1", "relu2_1", "relu3_1", "relu4_1", "relu5_1")
style_weight = 1e4
def get_style_op(encoder, layer_weight):
return ops.GramOperator(encoder, score_weight=layer_weight)
style_loss = ops.MultiLayerEncodingOperator(
multi_layer_encoder,
style_layers,
get_style_op,
score_weight=style_weight,
)
criterion = loss.PerceptualLoss(content_loss, style_loss).to(device)
print(criterion)
########################################################################################
# We set the target images for the optimization ``criterion``.
criterion.set_content_image(content_image)
criterion.set_style_image(style_image)
########################################################################################
# We perform the unguided NST and show the result.
score_weight=content_weight)
class GramOperator(ops.GramOperator):
def enc_to_repr(self, enc: torch.Tensor) -> torch.Tensor:
repr = super().enc_to_repr(enc)
num_channels = repr.size()[1]
return repr / num_channels
style_layers = ("relu1_2", "relu2_2", "relu3_3", "relu4_3")
style_weight = 1e10
style_loss = ops.MultiLayerEncodingOperator(
multi_layer_encoder,
style_layers,
lambda encoder, layer_weight: GramOperator(encoder,
score_weight=layer_weight),
layer_weights="sum",
score_weight=style_weight,
)
criterion = loss.PerceptualLoss(content_loss, style_loss)
criterion = criterion.to(device)
print(criterion)
########################################################################################
# Training
# --------
#
# In a first step we load the style image that will be used to train the
# ``transformer``.
