examples/tensorflow/benchmarking/plot_csv_file.py/Plot/plot
class Plot:    def plot(self):
        fig, ax = plt.subplots()
        title_str = "Time usage" if self.args.is_time else "Memory usage"
        title_str = title_str + " for training" if self.args.is_train else title_str + " for inference"

        if not self.args.no_log_scale:
            # set logarithm scales
            ax.set_xscale("log")
            ax.set_yscale("log")

        for axis in [ax.xaxis, ax.yaxis]:
            axis.set_major_formatter(ScalarFormatter())

        for model_name_idx, model_name in enumerate(self.result_dict.keys()):
            batch_sizes = sorted(set(self.result_dict[model_name]["bsz"]))
            sequence_lengths = sorted(set(self.result_dict[model_name]["seq_len"]))
            results = self.result_dict[model_name]["result"]

            (x_axis_array, inner_loop_array) = (
                (batch_sizes, sequence_lengths) if self.args.plot_along_batch else (sequence_lengths, batch_sizes)
            )

            label_model_name = (
                model_name if self.args.short_model_names is None else self.args.short_model_names[model_name_idx]
            )

            for inner_loop_value in inner_loop_array:
                if self.args.plot_along_batch:
                    y_axis_array = np.asarray(
                        [results[(x, inner_loop_value)] for x in x_axis_array if (x, inner_loop_value) in results],
                        dtype=int,
                    )
                else:
                    y_axis_array = np.asarray(
                        [results[(inner_loop_value, x)] for x in x_axis_array if (inner_loop_value, x) in results],
                        dtype=np.float32,
                    )

                (x_axis_label, inner_loop_label) = (
                    ("batch_size", "len") if self.args.plot_along_batch else ("in #tokens", "bsz")
                )

                x_axis_array = np.asarray(x_axis_array, int)[: len(y_axis_array)]
                plt.scatter(
                    x_axis_array, y_axis_array, label=f"{label_model_name} - {inner_loop_label}: {inner_loop_value}"
                )
                plt.plot(x_axis_array, y_axis_array, "--")

            title_str += f" {label_model_name} vs."

        title_str = title_str[:-4]
        y_axis_label = "Time in s" if self.args.is_time else "Memory in MB"

        # plot
        plt.title(title_str)
        plt.xlabel(x_axis_label)
        plt.ylabel(y_axis_label)
        plt.legend()

        if self.args.figure_png_file is not None:
            plt.savefig(self.args.figure_png_file)
        else:
            plt.show()