提交 6cbcadeb 编写于 作者: C chenfeiyu 提交者: liuyibing01

add layer_tools

上级 155dfe63
import numpy as np
from torch import nn
import paddle.fluid.dygraph as dg
def summary(layer):
num_params = num_elements = 0
print("layer summary:")
for name, param in layer.state_dict().items():
print("{}|{}|{}".format(name, param.shape, np.prod(param.shape)))
num_elements += np.prod(param.shape)
num_params += 1
print("layer has {} parameters, {} elements.".format(
num_params, num_elements))
def freeze(layer):
for param in layer.parameters():
param.trainable = False
def unfreeze(layer):
for param in layer.parameters():
param.trainable = True
def torch_summary(layer):
num_params = num_elements = 0
print("layer summary:")
for name, param in layer.named_parameters():
print("{}|{}|{}".format(name, param.shape, np.prod(param.shape)))
num_elements += np.prod(param.shape)
num_params += 1
print("layer has {} parameters, {} elements.".format(
num_params, num_elements))
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册