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.