ViT模型前面已定义
parameters
v = ViT(
image_size=224, # 输入图像的大小
patch_size=16, # 每个token/patch的大小16x16
num_classes=1000, # 多分类
dim=1024, # encoder规定的输入的维度
depth=6, # Encoder的个数
heads=16, # 多头注意力机制的head个数
mlp_dim=2048, # mlp的维度
dropout=0.1, #
emb_dropout=0.1 # embedding一半会接一个dropout
)
params = list(v.parameters())
k = 0
for i in params:
l = 1
print("该层的结构:" + str(list(i.size())))
for j in i.size():
l *= j
print("该层参数和:" + str(l))
k = k + l
print("总参数数量和:" + str(k))
得到params这个参数列表
对每一层的参数列表(如:c,w,h)打印出来
任何求乘积,就是每层的参数总数(如:c*w*h)
torchsummary
from torchsummary import summary
if __name__ == '__main__':
v = ViT(
image_size=224, # 输入图像的大小
patch_size=16, # 每个token/patch的大小16x16
num_classes=1000, # 多分类
dim=1024, # encoder规定的输入的维度
depth=6, # Encoder的个数
heads=16, # 多头注意力机制的head个数
mlp_dim=2048, # mlp的维度
dropout=0.1, #
emb_dropout=0.1 # embedding一半会接一个dropout
)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
summary(v.to(device), input_size=(3, 224, 224), batch_size=-1)
stat
from torchstat import stat
if __name__ == '__main__':
v = ViT(
image_size=224, # 输入图像的大小
patch_size=16, # 每个token/patch的大小16x16
num_classes=1000, # 多分类
dim=1024, # encoder规定的输入的维度
depth=6, # Encoder的个数
heads=16, # 多头注意力机制的head个数
mlp_dim=2048, # mlp的维度
dropout=0.1, #
emb_dropout=0.1 # embedding一半会接一个dropout
)
stat(v, (3, 224, 224))
shape只能是3维的。transformer的形状多变,且初始是4维。只能应用于某些模型
如alexnet。alexnet之前已定义
from torchstat import stat
net = AlexNet()
stat(net,(1,224,224))
thop
import torch
from thop import profile
if __name__ == '__main__':
v = ViT(
image_size=224, # 输入图像的大小
patch_size=16, # 每个token/patch的大小16x16
num_classes=1000, # 多分类
dim=1024, # encoder规定的输入的维度
depth=6, # Encoder的个数
heads=16, # 多头注意力机制的head个数
mlp_dim=2048, # mlp的维度
dropout=0.1, #
emb_dropout=0.1 # embedding一半会接一个dropout
)
input = torch.randn(1, 3, 224,224 )
flops, params = profile(v, inputs=(input,))
print('flops:{}'.format(flops))
print('params:{}'.format(params))