Allow feature_shape to be passed to fx plan initialize

[psfoley]

This PR:

• Fixes the feature_shape argument that can be passed to fx plan initialize. By passing this argument (which allows for list types arguments, such as [1,28,28]), loading the data loader for the model owner / aggregator can be avoided. This is important because the aggregator may not have access to local data to generate the initial model.

[MasterSkepticista]

While this method works for single-input models with static shapes, it may not work for multi-tensor input models, or single-input models with custom input dtypes.

A general approach could be for the aggregator to take InputSpec as an argument. InputSpec is could be a dictionary of list of TensorSpec[TensorShape, DType].

Example multi-input-output model:

# Can generalize if `ast` parses it as a dict instead
--input_shape {'input_0': [1, 240, 240, 4], 'output_1': [1, 240, 240, 1]}

# Assumes first input to take the given shape, if `ast` parses as list.
--input_shape [1, 240, 240, 1]

[psfoley]

While this method works for single-input models with static shapes, it may not work for multi-tensor input models, or single-input models with custom input dtypes.

A general approach could be for the aggregator to take InputSpec as an argument. InputSpec is could be a dictionary of list of TensorSpec[TensorShape, DType].

Example multi-input-output model:

# Can generalize if `ast` parses it as a dict instead
--input_shape {'input_0': [1, 240, 240, 4], 'output_1': [1, 240, 240, 1]}

# Assumes first input to take the given shape, if `ast` parses as list.
--input_shape [1, 240, 240, 1]

This is a good point. I tested that the existing code (that makes use of ast.literal_eval) can handle dictionaries as well. The main caveat is that the dictionary needs to be wrapped in quotes to be passed via command line. I've tried to make it clear in the help string that this is a requirement, and provided your example (wrapped in double quotes) as a reference for users.

[kta-intel]

This looks good to me, overall. It successfully avoids the need of loading the data at the aggregator to initialize the model, which is a great addition.

One comment: this won't always guarantee that the model is initialized using the specified input values or that the dataloader will load in data with the specified shape at the collaborators. For example, the torch_cnn_mnist workspace hard codes the model parameters in src/taskrunner.py so it will ultimately ignore the values passed through fx plan initialize --input_shape and initialize regardless, and even if we use .get_feature_shape() during initialization, only the channel size matters since the HxW isn't required in for torch conv layers.

In my opinion, this is not a breaking issue and this PR is good to go. However, going forward we should consider ways to reduce any potential confusion caused by this from a user perspective, such as having the existing workspace templates erroring out if an invalid input_shape is provided or issuing a user warning if it is ignored.

[teoparvanov]

Looks great, thanks @psfoley !

[teoparvanov]

I suppose you don't extend DataLoader, because only the get_feature_shape() method is required during plan initialization? In this case, would it make sense to call this class FeatureShapeLoader?