show non trainable params in summary, use pandas to display engineering units

This shows a trainable parameter in the full data frame and trainable vs non-trainable params in the summary. This should be like keras [example](https://github.com/keras-team/keras/pull/11795#issue-235955201).

I also changed the pandas display so that instead of putting (M) or (K) manually in the label, pandas appends it to the number automatically.
This commit is contained in:
Mike Clark
2019-07-07 03:35:52 +00:00
committed by GitHub
parent 87fa9ede4d
commit a903488e39
+30 -17
View File
@@ -32,9 +32,10 @@ def summary(model, x, *args, **kwargs):
info["ksize"] = "-"
info["inner"] = OrderedDict()
info["params"], info["macs"] = 0, 0
info["params_nt"], info["params"], info["macs"] = 0, 0, 0
for name, param in module.named_parameters():
info["params"] += param.nelement() * param.requires_grad
info["params_nt"] += param.nelement() * (not param.requires_grad)
if name == "weight":
ksize = list(param.size())
@@ -76,33 +77,45 @@ def summary(model, x, *args, **kwargs):
summary = OrderedDict()
model.apply(register_hook)
with torch.no_grad():
model(x) if not (kwargs or args) else model(x, *args, **kwargs)
for hook in hooks:
hook.remove()
try:
with torch.no_grad():
model(x) if not (kwargs or args) else model(x, *args, **kwargs)
finally:
for hook in hooks:
hook.remove()
# Use pandas to align the columns
df = pd.DataFrame(summary).T
df["Mult-Adds (M)"] = pd.to_numeric(df["macs"], errors="coerce")/1e6
df["Params (K)"] = pd.to_numeric(df["params"], errors="coerce")/1e3
df["Mult-Adds"] = pd.to_numeric(df["macs"], errors="coerce")
df["Params"] = pd.to_numeric(df["params"], errors="coerce")
df["Non-trainable params"] = pd.to_numeric(df["params_nt"], errors="coerce")
df = df.rename(columns=dict(
ksize="Kernel Shape",
out="Output Shape",
))
df.index.name = "Layer"
df = df[["Kernel Shape", "Output Shape", "Params (K)", "Mult-Adds (M)"]]
df_sum = df.sum()
df.index.name = "Layer"
df = df[["Kernel Shape", "Output Shape", "Params", "Mult-Adds"]]
max_repr_width = max([len(row) for row in df.to_string().split("\n")])
print("="*max_repr_width)
print(df.replace(np.nan, "-"))
print("-"*max_repr_width)
print("Params (K): ", df_sum["Params (K)"])
print("Mult-Adds (M): ", df_sum["Mult-Adds (M)"])
print("="*max_repr_width)
with pd.option_context("display.max_rows", 10, "display.max_columns", 5, 'display.float_format', pd.io.formats.format.EngFormatter(use_eng_prefix=True)):
print("="*max_repr_width)
print(df.replace(np.nan, "-"))
print("-"*max_repr_width)
df_total = pd.DataFrame(
{"Total params": (df_sum["Params"] + df_sum["params_nt"]),
"Trainable params": df_sum["Params"],
"Non-trainable params": df_sum["params_nt"],
"Mult-Adds": df_sum["Mult-Adds"]
},
index=['Totals']
).T
print(df_total)
print("="*max_repr_width)
return df