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")