def _resolve_transforms( self, running_stage: RunningStage) -> Optional[Dict[str, Callable]]: from flash.core.data.data_pipeline import DataPipeline resolved_function = getattr( self, DataPipeline._resolve_function_hierarchy("default_transforms", self, running_stage, Preprocess)) with CurrentRunningStageFuncContext(running_stage, "default_transforms", self): transforms: Optional[Dict[str, Callable]] = resolved_function() return transforms
def running_stage(self, running_stage: RunningStage) -> None: from flash.core.data.data_pipeline import DataPipeline # noqa F811 from flash.core.data.data_source import DataSource # noqa F811 # TODO: something better than this self._running_stage = running_stage self._load_sample_context = CurrentRunningStageFuncContext( self.running_stage, "load_sample", self.data_source) self.load_sample: Callable[[DATA_TYPE, Optional[Any]], Any] = getattr( self.data_source, DataPipeline._resolve_function_hierarchy( 'load_sample', self.data_source, self.running_stage, DataSource, ))
def generate_dataset( self, data: Optional[DATA_TYPE], running_stage: RunningStage, ) -> Optional[Union[AutoDataset, IterableAutoDataset]]: """Generate a single dataset with the given input to :meth:`~flash.core.data.data_source.DataSource.load_data` for the given ``running_stage``. Args: data: The input to :meth:`~flash.core.data.data_source.DataSource.load_data` to use to create the dataset. running_stage: The running_stage for this dataset. Returns: The constructed :class:`~flash.core.data.auto_dataset.BaseAutoDataset`. """ is_none = data is None if isinstance(data, Sequence): is_none = data[0] is None if not is_none: from flash.core.data.data_pipeline import DataPipeline mock_dataset = typing.cast(AutoDataset, MockDataset()) with CurrentRunningStageFuncContext(running_stage, "load_data", self): resolved_func_name = DataPipeline._resolve_function_hierarchy( "load_data", self, running_stage, DataSource) load_data: Callable[[DATA_TYPE, Optional[Any]], Any] = getattr(self, resolved_func_name) parameters = signature(load_data).parameters if len( parameters ) > 1 and "dataset" in parameters: # TODO: This was DATASET_KEY before data = load_data(data, mock_dataset) else: data = load_data(data) if has_len(data): dataset = AutoDataset(data, self, running_stage) else: dataset = IterableAutoDataset(data, self, running_stage) dataset.__dict__.update(mock_dataset.metadata) return dataset
def test_data_pipeline_is_overriden_and_resolve_function_hierarchy(tmpdir): class CustomPreprocess(DefaultPreprocess): def val_pre_tensor_transform(self, *_, **__): pass def predict_to_tensor_transform(self, *_, **__): pass def train_post_tensor_transform(self, *_, **__): pass def test_collate(self, *_, **__): pass def val_per_sample_transform_on_device(self, *_, **__): pass def train_per_batch_transform_on_device(self, *_, **__): pass def test_per_batch_transform_on_device(self, *_, **__): pass preprocess = CustomPreprocess() data_pipeline = DataPipeline(preprocess=preprocess) train_func_names: Dict[str, str] = { k: data_pipeline._resolve_function_hierarchy( k, data_pipeline._preprocess_pipeline, RunningStage.TRAINING, Preprocess) for k in data_pipeline.PREPROCESS_FUNCS } val_func_names: Dict[str, str] = { k: data_pipeline._resolve_function_hierarchy( k, data_pipeline._preprocess_pipeline, RunningStage.VALIDATING, Preprocess) for k in data_pipeline.PREPROCESS_FUNCS } test_func_names: Dict[str, str] = { k: data_pipeline._resolve_function_hierarchy( k, data_pipeline._preprocess_pipeline, RunningStage.TESTING, Preprocess) for k in data_pipeline.PREPROCESS_FUNCS } predict_func_names: Dict[str, str] = { k: data_pipeline._resolve_function_hierarchy( k, data_pipeline._preprocess_pipeline, RunningStage.PREDICTING, Preprocess) for k in data_pipeline.PREPROCESS_FUNCS } # pre_tensor_transform assert train_func_names["pre_tensor_transform"] == "pre_tensor_transform" assert val_func_names["pre_tensor_transform"] == "val_pre_tensor_transform" assert test_func_names["pre_tensor_transform"] == "pre_tensor_transform" assert predict_func_names["pre_tensor_transform"] == "pre_tensor_transform" # to_tensor_transform assert train_func_names["to_tensor_transform"] == "to_tensor_transform" assert val_func_names["to_tensor_transform"] == "to_tensor_transform" assert test_func_names["to_tensor_transform"] == "to_tensor_transform" assert predict_func_names[ "to_tensor_transform"] == "predict_to_tensor_transform" # post_tensor_transform assert train_func_names[ "post_tensor_transform"] == "train_post_tensor_transform" assert val_func_names["post_tensor_transform"] == "post_tensor_transform" assert test_func_names["post_tensor_transform"] == "post_tensor_transform" assert predict_func_names[ "post_tensor_transform"] == "post_tensor_transform" # collate assert train_func_names["collate"] == "collate" assert val_func_names["collate"] == "collate" assert test_func_names["collate"] == "test_collate" assert predict_func_names["collate"] == "collate" # per_sample_transform_on_device assert train_func_names[ "per_sample_transform_on_device"] == "per_sample_transform_on_device" assert val_func_names[ "per_sample_transform_on_device"] == "val_per_sample_transform_on_device" assert test_func_names[ "per_sample_transform_on_device"] == "per_sample_transform_on_device" assert predict_func_names[ "per_sample_transform_on_device"] == "per_sample_transform_on_device" # per_batch_transform_on_device assert train_func_names[ "per_batch_transform_on_device"] == "train_per_batch_transform_on_device" assert val_func_names[ "per_batch_transform_on_device"] == "per_batch_transform_on_device" assert test_func_names[ "per_batch_transform_on_device"] == "test_per_batch_transform_on_device" assert predict_func_names[ "per_batch_transform_on_device"] == "per_batch_transform_on_device" train_worker_preprocessor = data_pipeline.worker_preprocessor( RunningStage.TRAINING) val_worker_preprocessor = data_pipeline.worker_preprocessor( RunningStage.VALIDATING) test_worker_preprocessor = data_pipeline.worker_preprocessor( RunningStage.TESTING) predict_worker_preprocessor = data_pipeline.worker_preprocessor( RunningStage.PREDICTING) _seq = train_worker_preprocessor.per_sample_transform assert _seq.pre_tensor_transform.func == preprocess.pre_tensor_transform assert _seq.to_tensor_transform.func == preprocess.to_tensor_transform assert _seq.post_tensor_transform.func == preprocess.train_post_tensor_transform assert train_worker_preprocessor.collate_fn.func == preprocess.collate assert train_worker_preprocessor.per_batch_transform.func == preprocess.per_batch_transform _seq = val_worker_preprocessor.per_sample_transform assert _seq.pre_tensor_transform.func == preprocess.val_pre_tensor_transform assert _seq.to_tensor_transform.func == preprocess.to_tensor_transform assert _seq.post_tensor_transform.func == preprocess.post_tensor_transform assert val_worker_preprocessor.collate_fn.func == DataPipeline._identity assert val_worker_preprocessor.per_batch_transform.func == preprocess.per_batch_transform _seq = test_worker_preprocessor.per_sample_transform assert _seq.pre_tensor_transform.func == preprocess.pre_tensor_transform assert _seq.to_tensor_transform.func == preprocess.to_tensor_transform assert _seq.post_tensor_transform.func == preprocess.post_tensor_transform assert test_worker_preprocessor.collate_fn.func == preprocess.test_collate assert test_worker_preprocessor.per_batch_transform.func == preprocess.per_batch_transform _seq = predict_worker_preprocessor.per_sample_transform assert _seq.pre_tensor_transform.func == preprocess.pre_tensor_transform assert _seq.to_tensor_transform.func == preprocess.predict_to_tensor_transform assert _seq.post_tensor_transform.func == preprocess.post_tensor_transform assert predict_worker_preprocessor.collate_fn.func == preprocess.collate assert predict_worker_preprocessor.per_batch_transform.func == preprocess.per_batch_transform
def __resolve_transforms( self, running_stage: RunningStage) -> Optional[Dict[str, Callable]]: from flash.core.data.data_pipeline import DataPipeline transforms_out = {} stage = _STAGES_PREFIX[running_stage] # iterate over all transforms hook name for transform_name in InputTransformPlacement: transforms = {} transform_name = transform_name.value # iterate over all prefixes for key in ApplyToKeyPrefix: # get the resolved hook name based on the current stage resolved_name = DataPipeline._resolve_function_hierarchy( transform_name, self, running_stage, InputTransform) # check if the hook name is specialized is_specialized_name = resolved_name.startswith(stage) # get the resolved hook name for apply to key on the current stage resolved_apply_to_key_name = DataPipeline._resolve_function_hierarchy( f"{key}_{transform_name}", self, running_stage, InputTransform) # check if resolved hook name for apply to key is specialized is_specialized_apply_to_key_name = resolved_apply_to_key_name.startswith( stage) # check if they are overridden by the user resolve_name_overridden = DataPipeline._is_overridden( resolved_name, self, InputTransform) resolved_apply_to_key_name_overridden = DataPipeline._is_overridden( resolved_apply_to_key_name, self, InputTransform) if resolve_name_overridden and resolved_apply_to_key_name_overridden: # if both are specialized or both aren't specialized, raise a exception # It means there is priority to specialize hooks name. if not (is_specialized_name ^ is_specialized_apply_to_key_name): raise MisconfigurationException( f"Only one of {resolved_name} or {resolved_apply_to_key_name} can be overridden." ) method_name = resolved_name if is_specialized_name else resolved_apply_to_key_name else: method_name = resolved_apply_to_key_name if resolved_apply_to_key_name_overridden else resolved_name # get associated transform try: fn = getattr(self, method_name)() except AttributeError as e: raise AttributeError( str(e) + ". Hint: Call super().__init__(...) after setting all attributes." ) if not callable(fn): raise MisconfigurationException( f"The hook {method_name} should return a function.") # if the default hook is used, it should return identity, skip it. if fn is self._identity: continue # wrap apply to key hook into `ApplyToKeys` with the associated key. if method_name == resolved_apply_to_key_name: fn = ApplyToKeys(key.value, fn) if method_name not in transforms: transforms[method_name] = fn # store the transforms. if transforms: transforms = list(transforms.values()) transforms_out[transform_name] = Compose( transforms) if len(transforms) > 1 else transforms[0] return transforms_out