|
import torch |
|
import pandas as pd |
|
|
|
|
|
|
|
|
|
|
|
pd.set_option('display.float_format', '{:.4f}'.format) |
|
|
|
def get_stats(tensor, axis=None, label=None): |
|
if label is None: |
|
prefix = "" |
|
else: |
|
prefix = f"{label}_" |
|
|
|
if axis is None: |
|
return { |
|
f"{prefix}count": tensor.numel(), |
|
f"{prefix}avg": tensor.mean().item(), |
|
f"{prefix}min": tensor.min().item(), |
|
f"{prefix}med": tensor.median().item(), |
|
f"{prefix}max": tensor.max().item(), |
|
} |
|
else: |
|
return { |
|
f"{prefix}avg": tensor.mean(dim=axis), |
|
f"{prefix}min": tensor.min(dim=axis), |
|
f"{prefix}med": tensor.median(dim=axis), |
|
f"{prefix}max": tensor.max(dim=axis), |
|
} |
|
|
|
class SparseBlob: |
|
def __init__(self, blob_path): |
|
self.blob_path = blob_path |
|
d = torch.load(blob_path) |
|
self._rpt = d['df'] |
|
self._blob = d['blob'] |
|
|
|
@classmethod |
|
def print_help(cls): |
|
print("\n- Help ------------------") |
|
print("\nh = SparseBlob(\"path to sparsity blob\")\n") |
|
|
|
methods = [ |
|
'preview', 'ls_layers', |
|
'get_sparsity_by_short_id', |
|
'get_sparsity_by_row_id', |
|
'get_sparsity_of_tile', |
|
'show_help' |
|
] |
|
for method in methods: |
|
print(f'{cls.__name__}.{method}:{cls.__dict__[method].__doc__}\n') |
|
print("- End of Help ------------------") |
|
|
|
def show_help(self): |
|
""" |
|
print help for available function of SparseBlob |
|
eg. h.show_help() |
|
""" |
|
self.print_help() |
|
|
|
def preview(self): |
|
""" |
|
preview sparsity dataframe, intend to show row id, short id for look up |
|
eg. h.preview() |
|
""" |
|
pd.set_option('display.max_rows', None) |
|
print(self._rpt) |
|
pd.reset_option('display.max_rows') |
|
|
|
def ls_layers(self): |
|
""" |
|
list all available layer ids for look up |
|
eg. h.ls_layers() |
|
""" |
|
pd.set_option('display.max_rows', None) |
|
print(self._rpt.short_id) |
|
pd.reset_option('display.max_rows') |
|
|
|
def get_sparsity_by_short_id(self, short_id): |
|
""" |
|
return a sparsity stats of a layer via short_id lookup. |
|
eg. h.get_sparsity_by_short_id('tx.0.attn.v') |
|
""" |
|
return self._rpt[self._rpt.short_id == short_id].iloc[0] |
|
|
|
def get_sparsity_by_row_id(self, id): |
|
""" |
|
return a sparsity stats of a layer via row id lookup. |
|
eg. h.get_sparsity_by_row_id(36) |
|
""" |
|
return self._rpt.iloc[id] |
|
|
|
def get_sparsity_of_tile(self, lut_id, tile_id): |
|
""" |
|
zoom into a specific layer and a specific tile, |
|
return the sparsity stats of the tile down to col, row granularity |
|
eg. h.get_sparsity_by_row_id(36, (5, 6)) |
|
""" |
|
if isinstance(lut_id, int): |
|
row = self.get_sparsity_by_row_id(lut_id) |
|
elif isinstance(lut_id, str): |
|
row = self.get_sparsity_by_layer_id(lut_id) |
|
else: |
|
print("Invalid lookup id, use row number or short_id; you can do .preview() to find out.") |
|
|
|
l = row['layer_id'] |
|
max_nrow, max_ncol = self._blob[l]['tile_sparsity'].shape |
|
|
|
if not isinstance(tile_id, tuple) or len(tile_id) != 2: |
|
print(f"[Error] tile_id must be a tuple, eg .get_sparsity_of_tile({lut_id}, ( 5, 10))") |
|
else: |
|
if tile_id[0] >= max_nrow or tile_id[1] >= max_ncol: |
|
print(f"[Error] tile not exist: tile_id must be (0-{max_nrow-1}, 0-{max_ncol-1})") |
|
else: |
|
outdict={ |
|
"tile_id": tile_id, |
|
"layer_id": l, |
|
"layer_sparsity": row['sparsity'], |
|
"tiled by": row['tile_shape'], |
|
"tile_sparsity": self._blob[l]['tile_sparsity'][tile_id].item(), |
|
**get_stats(self._blob[l]['tile_sparsity_per_col'][tile_id], label="col"), |
|
**get_stats(self._blob[l]['tile_sparsity_per_row'][tile_id], label="row"), |
|
} |
|
self._print_one_tile_stats(outdict) |
|
|
|
def _print_one_tile_stats(self, d): |
|
for key, value in d.items(): |
|
if isinstance(value, float): |
|
print(f'{value:>40.4f} : {key}') |
|
else: |
|
print(f'{str(value):>40} : {key}') |
|
|
|
|
|
SparseBlob.print_help() |
