def print_parameters(model):
# compute number of parameters
num_params, _ = get_num_parameter(model, trainable=False)
num_bytes = num_params * 32 // 8 # assume float32 for all
print(
f"Number of parameters: {human_format(num_params)} ({sizeof_fmt(num_bytes)} for float32)"
)
num_trainable_params, trainable_parameters = get_num_parameter(
model, trainable=True)
print("Number of trainable parameters:",
human_format(num_trainable_params))
if config["only_list_parameters"]:
# Print detailed number of parameters
print(tabulate.tabulate(trainable_parameters))
def print_flops(model):
shape = None
if config["dataset"] in ["Cifar10", "Cifar100"]:
shape = (1, 3, 32, 32)
else:
print(
f"Unknown dataset {config['dataset']} input size to compute # FLOPS"
)
return
try:
from thop import profile
except BaseException:
print("Please `pip install thop` to compute # FLOPS")
return
model = model.train()
input_data = torch.rand(*shape)
num_flops, num_params = profile(model, inputs=(input_data, ))
print("Number of FLOPS:", human_format(num_flops))