def forward(self, input: Dict[str, Tensor]) -> Dict[str, Tensor]:
targets = input.keys()
operation_outputs = [op(input) for op in self.operations]
output = {}
for target in targets:
with torch.set_grad_enabled(target_requires_grad(target)):
output[target] = (
torch.stack([op[target] for op in operation_outputs]) * self.weights.view(-1, 1, 1, 1, 1)
).sum(0)
return output
def forward(self, input):
mask = self.get_probability_mask(input["image_batch"].size(0))
operation_output = self.operation(input)
targets = input.keys()
output = {}
for target in targets:
with torch.set_grad_enabled(target_requires_grad(target)):
output[target] = (mask * operation_output[target] +
(1 - mask) * input[target]).clamp_(0, 1)
return output
def operation(self, input):
magnitude = self.magnitude
value = convert_value_range(
magnitude, self.value_range) if magnitude is not None else None
output = {}
targets = input.keys()
for target in targets:
operation_fn = self.targets[target]
operation_output = operation_fn(input[target], value)
if target_requires_grad(target) and self.ste:
operation_output = ste(operation_output, magnitude)
output[target] = operation_output
return output