python-计算trocr模型的参数量（Mb)

import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

def count_model_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def parameters_size_in_MB(model):
    param_count = count_model_parameters(model)
    param_size_MB = param_count * 4 / (1024 ** 2)  # 4 bytes per parameter (float32)
    return param_size_MB

# Load TrOCR model
model_name = "microsoft/trocr-base-stage1"
model = VisionEncoderDecoderModel.from_pretrained(model_name)

# Calculate the size of the parameters in MB
param_size = count_model_parameters(model)
param_size_MB = parameters_size_in_MB(model)
print(f"Size of parameters in {model_name}: {param_size}  {param_size_MB:.2f} MB")