Flow Generator

Extensions that are used for training the flow generator.

Accessor: flow_generator

class animal_soup.extensions.FlowGeneratorDataframeExtension(df)

Pandas dataframe extensions for training the flow generator.

train(mode='slow', flow_window=11, batch_size=32, gpu_id=0, initial_lr=0.0001, stop_method='learning_rate', model_in=None, model_out=None)

Train flow generator model.

Parameters:

• mode (str, default 'slow') –

  Argument must be one of [“slow”, “medium”, “fast”]. Determines the model used for training the flow generator.

  mode | model
  fast	TinyMotionNet
  medium	MotionNet
  slow	TinyMotionNet3D

• flow_window (int, default 11) – Window size for computing optical flow features of a frame. Recommended to be the minimum number of frames any given behavior takes.

• batch_size (int, default 32) – Batch size.

• gpu_id (int, default 0) – Specify which gpu to use for training the model.

• initial_lr (float, default 0.0001) – Initial learning rate.

• stop_method (str, default learning_rate) –

  Method for stopping training. Argument must be one of [“learning_rate”, “num_epochs”]

  stop method	description
  learning_rate	Stop training when learning rate drops below a given threshold, means loss has stopped improving
  num_epochs	Stop training after a given number of epochs

• model_in (str or Path, default None) – If you want to retrain the model using different model weights than the default. User can provide a location to a different model checkpoint. For example, if you had retrained the flow generator previously and wanted to use those weights instead.

• model_out (str or Path, default None) – User provided location of where to store model output such as model checkpoint with updated weights, hdf5 file with model results/metrics, etc. Should be a directory. By default, the model output will get stored in the same directory as the dataframe.