huggingface | how to get the params‘ number of a model

Posted 2023-03-02 CSU迦叶

主要实现在下列的how_big(model)函数中

# 探究plbart和codegpt的参数量
from transformers import PLBartModel,GPT2Model


def how_big(model):
    total_num = sum(p.numel() for p in model.parameters())
    trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print("Total: ,Trainable: ".format(total_num,trainable_num))


def main():
    bart = PLBartModel.from_pretrained("uclanlp/plbart-java-cs")
    how_big(bart)
    codegpt = GPT2Model.from_pretrained("pytorch.bin所在的文件夹路径")
    how_big(codegpt)

if __name__ == "__main__":
    main()

输出结果为

Total: 139220736,Trainable: 139220736
Total: 124442112,Trainable: 124442112