91 lines
3.4 KiB
Python

import pandas as pd
pd.set_option('display.width', 1000)
pd.set_option('display.max_rows', 10000)
pd.set_option('display.max_columns', 10000)
def round_value(value, binary=False):
divisor = 1024. if binary else 1000.
if value // divisor**4 > 0:
return str(round(value / divisor**4, 2)) + 'T'
elif value // divisor**3 > 0:
return str(round(value / divisor**3, 2)) + 'G'
elif value // divisor**2 > 0:
return str(round(value / divisor**2, 2)) + 'M'
elif value // divisor > 0:
return str(round(value / divisor, 2)) + 'K'
return str(value)
def report_format(collected_nodes):
data = list()
for node in collected_nodes:
name = node.name
input_shape = ' '.join(
['{:>3d}'] *
len(node.input_shape)).format(*[e for e in node.input_shape])
output_shape = ' '.join(
['{:>3d}'] *
len(node.output_shape)).format(*[e for e in node.output_shape])
parameter_quantity = node.parameter_quantity
inference_memory = node.inference_memory
MAdd = node.MAdd
Flops = node.Flops
mread, mwrite = [i for i in node.Memory]
duration = node.duration
data.append([
name, input_shape, output_shape, parameter_quantity,
inference_memory, MAdd, duration, Flops, mread, mwrite
])
df = pd.DataFrame(data)
df.columns = [
'module name', 'input shape', 'output shape', 'params', 'memory(MB)',
'MAdd', 'duration', 'Flops', 'MemRead(B)', 'MemWrite(B)'
]
df['duration[%]'] = df['duration'] / (df['duration'].sum() + 1e-7)
df['MemR+W(B)'] = df['MemRead(B)'] + df['MemWrite(B)']
total_parameters_quantity = df['params'].sum()
total_memory = df['memory(MB)'].sum()
total_operation_quantity = df['MAdd'].sum()
total_flops = df['Flops'].sum()
total_duration = df['duration[%]'].sum()
total_mread = df['MemRead(B)'].sum()
total_mwrite = df['MemWrite(B)'].sum()
total_memrw = df['MemR+W(B)'].sum()
del df['duration']
# Add Total row
total_df = pd.Series([
total_parameters_quantity, total_memory, total_operation_quantity,
total_flops, total_duration, mread, mwrite, total_memrw
],
index=[
'params', 'memory(MB)', 'MAdd', 'Flops',
'duration[%]', 'MemRead(B)', 'MemWrite(B)',
'MemR+W(B)'
],
name='total')
df = df.append(total_df)
df = df.fillna(' ')
df['memory(MB)'] = df['memory(MB)'].apply(lambda x: '{:.2f}'.format(x))
df['duration[%]'] = df['duration[%]'].apply(lambda x: '{:.2%}'.format(x))
df['MAdd'] = df['MAdd'].apply(lambda x: '{:,}'.format(x))
df['Flops'] = df['Flops'].apply(lambda x: '{:,}'.format(x))
summary = str(df) + '\n'
summary += "=" * len(str(df).split('\n')[0])
summary += '\n'
summary += "Total params: {:,}\n".format(total_parameters_quantity)
summary += "-" * len(str(df).split('\n')[0])
summary += '\n'
summary += "Total memory: {:.2f}MB\n".format(total_memory)
summary += "Total MAdd: {}MAdd\n".format(
round_value(total_operation_quantity))
summary += "Total Flops: {}Flops\n".format(round_value(total_flops))
summary += "Total MemR+W: {}B\n".format(round_value(total_memrw, True))
return summary