def test_MultiOperatorLoss_trim(self):
class TestOperator(EncodingOperator):
def __init__(self, encoder, **kwargs):
super().__init__(**kwargs)
self._encoder = encoder
@property
def encoder(self):
return self._encoder
def process_input_image(self, image):
pass
layers = [str(idx) for idx in range(3)]
modules = [(layer, nn.Module()) for layer in layers]
multi_layer_encoder = MultiLayerEncoder(modules)
ops = ((
"op",
TestOperator(
multi_layer_encoder.extract_single_layer_encoder(layers[0])),
), )
loss.MultiOperatorLoss(ops, trim=True)
self.assertTrue(layers[0] in multi_layer_encoder)
for layer in layers[1:]:
self.assertFalse(layer in multi_layer_encoder)
def make_loss(
loss_str: str, layers_str: str, score_weight: float,
mle: enc.MultiLayerEncoder
) -> Union[loss.ComparisonLoss, loss.MultiLayerEncodingLoss]:
loss_str_normalized = loss_str.lower().replace("_", "").replace("-", "")
if loss_str_normalized not in LOSSES.keys():
raise ValueError(
add_suggestion(
f"Unknown loss '{loss_str}'.",
word=loss_str_normalized,
possibilities=tuple(zip(*LOSSES.values()))[0],
))
_, loss_fn = LOSSES[loss_str_normalized]
layers = [layer.strip() for layer in layers_str.split(",")]
layers = sorted(layer for layer in layers if layer)
for layer in layers:
mle.verify(layer)
if len(layers) == 1:
return loss_fn(mle.extract_encoder(layers[0]), score_weight)
else:
return loss.MultiLayerEncodingLoss(mle,
layers,
loss_fn,
score_weight=score_weight)
def test_MultiOperatorLoss_call_encode(self):
class TestOperator(EncodingOperator):
def __init__(self, encoder, **kwargs):
super().__init__(**kwargs)
self._encoder = encoder
@property
def encoder(self):
return self._encoder
def process_input_image(self, image):
return torch.sum(image)
count = ForwardPassCounter()
modules = (("count", count),)
multi_layer_encoder = MultiLayerEncoder(modules)
ops = [
(str(idx), TestOperator(multi_layer_encoder.extract_encoder("count")),)
for idx in range(3)
]
multi_op_loss = loss.MultiOperatorLoss(ops)
torch.manual_seed(0)
input = torch.rand(1, 3, 128, 128)
multi_op_loss(input)
actual = count.count
desired = 1
self.assertEqual(actual, desired)
multi_op_loss(input)
actual = count.count
desired = 2
self.assertEqual(actual, desired)
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)
def multi_layer_encoder_with_layer():
layer = "conv"
multi_layer_encoder = MultiLayerEncoder(((layer, nn.Conv2d(3, 3, 1)), ))
return multi_layer_encoder, layer