lambda fn: add prob and targets args for training #1034
Description
It's sometomes easier and more understandable to use lambda functions to perform some steps for data processings/augs. For example, flipping data along Z:
{
// imgaug doesn't flip the last dimension
"@type": "lambda",
lambda_str: "lambda x: x[..., ::-1]"
prob: 0.5,
}
The problem is that lambda doesn't natively support prob. Or targets for selective application.
As a prototype, I'm duplicating the lambda implementation with a "training" variant:
from zetta_utils.tensor_ops.common import supports_dict
from zetta_utils.augmentations import prob_aug
@register("lambda_training", allow_partial=False)
def efficient_parse_lambda_str_training(
lambda_str: str,
name: Optional[str] = None,
prob: float | None = None,
targets: Container[str] | None = None,
) -> Callable:
"""Parses strings that are lambda functions"""
if not isinstance(lambda_str, str):
raise TypeError("`lambda_str` must be a string.")
if not lambda_str.startswith("lambda"):
raise ValueError("`lambda_str` must start with 'lambda'.")
if len(lambda_str) > LAMBDA_STR_MAX_LENGTH:
raise ValueError(f"`lambda_str` must be at most {LAMBDA_STR_MAX_LENGTH} characters.")
return BuilderPartial(
spec={
"@type": "invoke_lambda_str_training",
"lambda_str": lambda_str,
"prob": prob,
"targets": targets,
},
name=name,
)
@register("invoke_lambda_str_training", allow_partial=False)
@prob_aug
@supports_dict
def invoke_lambda_str(*args: list, lambda_str: str, **kwargs: dict) -> Any:
return eval(lambda_str)(*args, **kwargs) # pylint: disable=eval-used
It's actually hard to directly modify and use one single lambda definition because of looping dependency with supports_dict and prob_aug. They are also not needed outside of training pipeline.