1 可视化模型结构以及每一层的shape和参数量。
-
安装 torchinfo 包
pip install torchinfo -i https://pypi.tuna.tsinghua.edu.cn/simple -
写代码
from torchinfo import summary
base_model = uniformer_small()
print(summary(base_model, (1, 3, 8, 224, 224)))
1 可视化模型结构以及每一层的shape和参数量。
安装 torchinfo 包
pip install torchinfo -i https://pypi.tuna.tsinghua.edu.cn/simple
写代码
from torchinfo import summary
base_model = uniformer_small()
print(summary(base_model, (1, 3, 8, 224, 224)))