model_benchmark.py

Save the number of trainable parameter and inference speed of all available models.

model_benchmark.benchmark(args: Namespace, device_handle) DataFrame[source]

Run the benchmark on all models.

Parameters:
argsNamespace

Arguments for configuring the benchmark.

Returns:
pd.DataFrame

A DataFrame with the benchmark results.

model_benchmark.estimate_inference_time(args: Namespace, model: BaseModel, input_size: Tuple[int, int], dtype_str: str) float[source]

Compute the average forward time for one model.

Parameters:
argsNamespace

Arguments for configuring the benchmark.

modelBaseModel

The model to perform the estimation.

Returns:
float

The average time of the runs.

model_benchmark.save_plot(output_dir: str | Path, model_name: str, df: DataFrame, plot_axes: Tuple[str, str] | None, log_x: bool, log_y: bool, datatype: str) None[source]

Create a plot of the results and save to disk.

Parameters:
output_dirUnion[str, Path]

Path to the directory where the plot will be saved.

model_namestr

Name of the model. Used just to name the resulting file.

dfpd.DataFrame

A DataFrame with the benchmark results.

plot_axesOptional[Tuple[str, str]]

Name of two parameters to create the scatter plot.

log_xbool

If set, the X-axis is plot in log scale.

log_ybool

If set, the Y-axis is plot in log scale.

datatypestr

Name of the datatype.