def generate_summary(
self,
include_macs: bool = True
) -> Sequence[Sequence[Union[str, float]]]:
summary = [
["Total params",
_number_as_readable_str(self.weight_count())],
[
"Trainable params",
_number_as_readable_str(self.weight_count(trainable=True)),
],
[
"Non-trainable params",
_number_as_readable_str(self.weight_count(trainable=False)),
],
["Model size", memory_as_readable_str(self.memory)],
[
"Model size (8-bit FP weights)",
memory_as_readable_str(self.int8_fp_weights_memory),
],
[
"Float-32 Equivalent",
memory_as_readable_str(self.fp_equivalent_memory)
],
[
"Compression Ratio of Memory",
self.memory / max(1e-8, self.fp_equivalent_memory),
],
]
if include_macs:
binarization_ratio = self.op_count("mac", 1) / max(
1, self.op_count(op_type="mac"))
ternarization_ratio = self.op_count("mac", 2) / max(
1, self.op_count(op_type="mac"))
summary.append([
"Number of MACs",
_number_as_readable_str(self.op_count(op_type="mac")),
])
if binarization_ratio > 0:
summary.append([
"Ratio of MACs that are binarized",
f"{binarization_ratio:.4f}"
])
if ternarization_ratio > 0:
summary.append([
"Ratio of MACs that are ternarized",
f"{ternarization_ratio:.4f}"
])
return summary
def test_memory_as_readable_str():
correct_strings = [ # 2^i bits, from i = 0 to 74
"0.12 B",
"0.25 B",
"0.50 B",
"1.00 B",
"2.00 B",
"4.00 B",
"8.00 B",
"16.00 B",
"32.00 B",
"64.00 B",
"128.00 B",
"256.00 B",
"512.00 B",
"1.00 KiB",
"2.00 KiB",
"4.00 KiB",
"8.00 KiB",
"16.00 KiB",
"32.00 KiB",
"64.00 KiB",
"128.00 KiB",
"256.00 KiB",
"512.00 KiB",
"1.00 MiB",
"2.00 MiB",
"4.00 MiB",
"8.00 MiB",
"16.00 MiB",
"32.00 MiB",
"64.00 MiB",
"128.00 MiB",
"256.00 MiB",
"512.00 MiB",
"1.00 GiB",
"2.00 GiB",
"4.00 GiB",
"8.00 GiB",
"16.00 GiB",
"32.00 GiB",
"64.00 GiB",
"128.00 GiB",
"256.00 GiB",
"512.00 GiB",
"1,024.00 GiB",
]
for i, correct_string in enumerate(correct_strings):
assert utils.memory_as_readable_str(2 ** i) == correct_string