armory
armory copied to clipboard
Create `fit` API with better user control
Currently, there are limited controls for training models in Armory.
Ideally, this would be something that Armory would call out to, provide a PyTorch, TF, or Numpy (or other) dataset, potentially with preprocessing applied, and enable the user full control of model training. They could then return a fit model. They could also use this to pass in various kwargs as desired from the config. Armory would still be in control of what dataset / split is passed in, but the remaining work could be left up to the user.
The current approach in Armory scenario code is essentially as follows:
def fit(self):
self.model.fit_generator(self.train_dataset, **self.fit_kwargs)
which uses the fit_generator()
API that ART adheres to. This doesn't seem too far off the description you've provided: Armory controls the data, user fit kwargs are passed in from the config, and users can define fit_generator()
as desired. What sort of additions/refactors/etc. are you envisioning would enhance user control?