plot_flatness.py

import os

import flatquant.utils as utils
import flatquant.args_utils as args_utils
import flatquant.flatness as flatness


def main():
    args, logger = args_utils.parser_gen()
    utils.seed_everything(seed=args.seed)
    args.vis_dir = os.path.join(args.output_dir, "flatness", args.model_name)
    os.makedirs(args.vis_dir, exist_ok=True)

    # get activation scales for smoothquant
    if not os.path.exists(f"./act_scales/{args.model_name}.pt"):
        flatness.get_act_scales(args, logger)

    # get pre-quantization transformations
    flatness_flatquant = flatness.get_flatness(args, logger, transform_type="flatquant")
    flatness_hadamard = flatness.get_flatness(args, logger, transform_type="hadamard")
    flatness_smoothquant = flatness.get_flatness(args, logger, transform_type="smoothquant")
    flatness_vanilla = flatness.get_flatness(args, logger, transform_type=None)
    
    # plot flatness
    flatness.plot_flatness_all_layers(args, flatness_flatquant, flatness_hadamard, flatness_smoothquant, flatness_vanilla)


if __name__ == '__main__':
    main()