def test_is_overridden_recursive(tmpdir): class TestInputTransform(InputTransform): @staticmethod def custom_transform(x): return x def collate(self): return self.custom_transform def val_collate(self): return self.custom_transform input_transform = TestInputTransform() assert DataPipeline._is_overridden_recursive("collate", input_transform, InputTransform, prefix="val") assert DataPipeline._is_overridden_recursive("collate", input_transform, InputTransform, prefix="train") assert not DataPipeline._is_overridden_recursive( "per_batch_transform_on_device", input_transform, InputTransform, prefix="train") assert not DataPipeline._is_overridden_recursive( "per_batch_transform_on_device", input_transform, InputTransform) with pytest.raises( MisconfigurationException, match="This function doesn't belong to the parent class"): assert not DataPipeline._is_overridden_recursive( "chocolate", input_transform, InputTransform)
def __configure_worker_and_device_collate_fn( running_stage: RunningStage, input_transform: InputTransform) -> Tuple[Callable, Callable]: from flash.core.data.data_pipeline import DataPipeline prefix: str = _STAGES_PREFIX[running_stage] transform_for_stage: _InputTransformPerStage = input_transform._transform[ running_stage] per_batch_transform_overridden: bool = DataPipeline._is_overridden_recursive( "per_batch_transform", input_transform, InputTransform, prefix=prefix) per_sample_transform_on_device_overridden: bool = DataPipeline._is_overridden_recursive( "per_sample_transform_on_device", input_transform, InputTransform, prefix=prefix) is_per_overridden = per_batch_transform_overridden and per_sample_transform_on_device_overridden if transform_for_stage.collate_in_worker_from_transform is None and is_per_overridden: raise MisconfigurationException( f"{input_transform.__class__.__name__}: `per_batch_transform` and `per_sample_transform_on_device` " f"are mutually exclusive for stage {running_stage}") if isinstance(transform_for_stage.collate_in_worker_from_transform, bool): worker_collate_fn, device_collate_fn = __make_collates( input_transform, not transform_for_stage.collate_in_worker_from_transform, input_transform._collate) else: worker_collate_fn, device_collate_fn = __make_collates( input_transform, per_sample_transform_on_device_overridden, input_transform._collate) worker_collate_fn = (worker_collate_fn.collate_fn if isinstance( worker_collate_fn, _InputTransformProcessor) else worker_collate_fn) return worker_collate_fn, device_collate_fn