File size: 1,204 Bytes
5ce8761 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import argparse
def str_none(value):
if value.lower() in ['none', 'null', 'nil'] or len(value) == 0:
return None
else:
return value
def str2bool(value):
if value.lower() in ['true', '1', 't', 'y', 'yes']:
return True
elif value.lower() in ['false', '0', 'f', 'n', 'no']:
return False
else:
raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")
def round_floats(o):
if isinstance(o, float): return round(o, 2)
if isinstance(o, dict): return {k: round_floats(v) for k, v in o.items()}
if isinstance(o, (list, tuple)): return [round_floats(x) for x in o]
return o
def count_parameters(model):
print(
"Trainable model parameters:",
sum(p.numel() for p in model.parameters() if p.requires_grad)
)
# Print number of trainable parameters for main modules
for name, submodule in model.named_modules():
if '.' not in name:
submodule_params = sum(
p.numel() for p in submodule.parameters()
if p.requires_grad
)
if submodule_params > 0:
print(f"{name} - trainable params: {submodule_params}")
|