计算 PyTorch 模型的参数量和 FLOPs
from mmcv.cnn import get_model_complexity_info
# from ptflops import get_model_complexity_info
flops, params = get_model_complexity_info(
model,
input_shape=(3, train_size, train_size), # input_res=(3, train_size, train_size)
)
print(params)
print(flops)
#下面是计算fps
bc = 1
rgb = torch.rand(bc, 3, train_size, train_size).cuda()
for i in range(50):
with torch.no_grad():
res = net(rgb)
frame_rate = np.zeros((1000, 1))
for i in range(1000):
with torch.no_grad():
start_time = time.time()
res = net(rgb)
end_time = time.time()
running_frame_rate = (1 * float((bc / (end_time - start_time))))
# print(i, '->', running_frame_rate)
frame_rate[i] = running_frame_rate
print(np.mean(frame_rate))
输出结果相同
-------------------------------------------mmcv.cnn---------------------------------------------
UNext(
1.472 M, 100.000% Params, 2.286 GFLOPs, 100.000% FLOPs,
(encoder1): Conv2d(0.0 M, 0.030% Params, 0.117 GFLOPs, 5.138% FLOPs, 3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(encoder2): Conv2d(0.005 M, 0.315% Params, 0.304 GFLOPs, 13.303% FLOPs, 16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(encoder3): Conv2d(0.037 M, 2.513% Params, 0.606 GFLOPs, 26.513% FLOPs, 32, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(ebn1): BatchNorm2d(0.0 M, 0.002% Params, 0.008 GFLOPs, 0.367% FLOPs, 16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(ebn2): BatchNorm2d(0.0 M, 0.004% Params, 0.004 GFLOPs, 0.183% FLOPs, 32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(ebn3): BatchNorm2d(0.0 M, 0.017% Params, 0.004 GFLOPs, 0.183% FLOPs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(norm3): LayerNorm(0.0 M, 0.022% Params, 0.0 GFLOPs, 0.014% FLOPs, (160,), eps=1e-05, elementwise_affine=True)
(norm4): LayerNorm(0.001 M, 0.035% Params, 0.0 GFLOPs, 0.006% FLOPs, (256,), eps=1e-05, elementwise_affine=True)
(dnorm3): LayerNorm(0.0 M, 0.022% Params, 0.0 GFLOPs, 0.014% FLOPs, (160,), eps=1e-05, elementwise_affine=True)
(dnorm4): LayerNorm(0.0 M, 0.017% Params, 0.001 GFLOPs, 0.046% FLOPs, (128,), eps=1e-05, elementwise_affine=True)
(block1): ModuleList(
(0): shiftedBlock(
0.053 M, 3.631% Params, 0.054 GFLOPs, 2.380% FLOPs,
(drop_path): Identity(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, )
(norm2): LayerNorm(0.0 M, 0.022% Params, 0.0 GFLOPs, 0.014% FLOPs, (160,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
0.053 M, 3.609% Params, 0.054 GFLOPs, 2.365% FLOPs,
(fc1): Linear(0.026 M, 1.750% Params, 0.026 GFLOPs, 1.147% FLOPs, in_features=160, out_features=160, bias=True)
(dwconv): DWConv(
0.002 M, 0.109% Params, 0.002 GFLOPs, 0.072% FLOPs,
(dwconv): Conv2d(0.002 M, 0.109% Params, 0.002 GFLOPs, 0.072% FLOPs, 160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=160)
)
(act): GELU(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, approximate='none')
(fc2): Linear(0.026 M, 1.750% Params, 0.026 GFLOPs, 1.147% FLOPs, in_features=160, out_features=160, bias=True)
(drop): Dropout(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, p=0.0, inplace=False)
)
)
)
(block2): ModuleList(
(0): shiftedBlock(
0.135 M, 9.148% Params, 0.034 GFLOPs, 1.502% FLOPs,
(drop_path): Identity(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, )
(norm2): LayerNorm(0.001 M, 0.035% Params, 0.0 GFLOPs, 0.006% FLOPs, (256,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
0.134 M, 9.113% Params, 0.034 GFLOPs, 1.497% FLOPs,
(fc1): Linear(0.066 M, 4.470% Params, 0.017 GFLOPs, 0.734% FLOPs, in_features=256, out_features=256, bias=True)
(dwconv): DWConv(
0.003 M, 0.174% Params, 0.001 GFLOPs, 0.029% FLOPs,
(dwconv): Conv2d(0.003 M, 0.174% Params, 0.001 GFLOPs, 0.029% FLOPs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=256)
)
(act): GELU(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, approximate='none')
(fc2): Linear(0.066 M, 4.470% Params, 0.017 GFLOPs, 0.734% FLOPs, in_features=256, out_features=256, bias=True)
(drop): Dropout(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, p=0.0, inplace=False)
)
)
)
(dblock1): ModuleList(
(0): shiftedBlock(
0.053 M, 3.631% Params, 0.054 GFLOPs, 2.380% FLOPs,
(drop_path): Identity(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, )
(norm2): LayerNorm(0.0 M, 0.022% Params, 0.0 GFLOPs, 0.014% FLOPs, (160,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
0.053 M, 3.609% Params, 0.054 GFLOPs, 2.365% FLOPs,
(fc1): Linear(0.026 M, 1.750% Params, 0.026 GFLOPs, 1.147% FLOPs, in_features=160, out_features=160, bias=True)
(dwconv): DWConv(
0.002 M, 0.109% Params, 0.002 GFLOPs, 0.072% FLOPs,
(dwconv): Conv2d(0.002 M, 0.109% Params, 0.002 GFLOPs, 0.072% FLOPs, 160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=160)
)
(act): GELU(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, approximate='none')
(fc2): Linear(0.026 M, 1.750% Params, 0.026 GFLOPs, 1.147% FLOPs, in_features=160, out_features=160, bias=True)
(drop): Dropout(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, p=0.0, inplace=False)
)
)
)
(dblock2): ModuleList(
(0): shiftedBlock(
0.035 M, 2.348% Params, 0.141 GFLOPs, 6.147% FLOPs,
(drop_path): Identity(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, )
(norm2): LayerNorm(0.0 M, 0.017% Params, 0.001 GFLOPs, 0.046% FLOPs, (128,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
0.034 M, 2.331% Params, 0.139 GFLOPs, 6.101% FLOPs,
(fc1): Linear(0.017 M, 1.122% Params, 0.067 GFLOPs, 2.936% FLOPs, in_features=128, out_features=128, bias=True)
(dwconv): DWConv(
0.001 M, 0.087% Params, 0.005 GFLOPs, 0.229% FLOPs,
(dwconv): Conv2d(0.001 M, 0.087% Params, 0.005 GFLOPs, 0.229% FLOPs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128)
)
(act): GELU(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, approximate='none')
(fc2): Linear(0.017 M, 1.122% Params, 0.067 GFLOPs, 2.936% FLOPs, in_features=128, out_features=128, bias=True)
(drop): Dropout(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, p=0.0, inplace=False)
)
)
)
(patch_embed3): OverlapPatchEmbed(
0.185 M, 12.555% Params, 0.189 GFLOPs, 8.278% FLOPs,
(proj): Conv2d(0.184 M, 12.533% Params, 0.189 GFLOPs, 8.264% FLOPs, 128, 160, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(norm): LayerNorm(0.0 M, 0.022% Params, 0.0 GFLOPs, 0.014% FLOPs, (160,), eps=1e-05, elementwise_affine=True)
)
(patch_embed4): OverlapPatchEmbed(
0.369 M, 25.097% Params, 0.095 GFLOPs, 4.137% FLOPs,
(proj): Conv2d(0.369 M, 25.062% Params, 0.094 GFLOPs, 4.131% FLOPs, 160, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(norm): LayerNorm(0.001 M, 0.035% Params, 0.0 GFLOPs, 0.006% FLOPs, (256,), eps=1e-05, elementwise_affine=True)
)
(decoder1): Conv2d(0.369 M, 25.055% Params, 0.094 GFLOPs, 4.130% FLOPs, 256, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder2): Conv2d(0.184 M, 12.531% Params, 0.189 GFLOPs, 8.263% FLOPs, 160, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder3): Conv2d(0.037 M, 2.507% Params, 0.151 GFLOPs, 6.611% FLOPs, 128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder4): Conv2d(0.005 M, 0.314% Params, 0.076 GFLOPs, 3.314% FLOPs, 32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder5): Conv2d(0.002 M, 0.158% Params, 0.152 GFLOPs, 6.651% FLOPs, 16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(dbn1): BatchNorm2d(0.0 M, 0.022% Params, 0.0 GFLOPs, 0.004% FLOPs, 160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(dbn2): BatchNorm2d(0.0 M, 0.017% Params, 0.0 GFLOPs, 0.011% FLOPs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(dbn3): BatchNorm2d(0.0 M, 0.004% Params, 0.0 GFLOPs, 0.011% FLOPs, 32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(dbn4): BatchNorm2d(0.0 M, 0.002% Params, 0.001 GFLOPs, 0.023% FLOPs, 16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(final): Conv2d(0.0 M, 0.002% Params, 0.009 GFLOPs, 0.390% FLOPs, 16, 2, kernel_size=(1, 1), stride=(1, 1))
(soft): Softmax(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, dim=1)
)
1.47 M
2.29 GFLOPs
frame_rate 188.31664432728067
-------------------------------------------ptflops----------------------------------------------
UNext(
1.47 M, 100.000% Params, 2.29 GMac, 98.851% MACs,
(encoder1): Conv2d(448, 0.030% Params, 117.44 MMac, 5.077% MACs, 3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(encoder2): Conv2d(4.64 k, 0.315% Params, 304.09 MMac, 13.146% MACs, 16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(encoder3): Conv2d(36.99 k, 2.513% Params, 606.08 MMac, 26.201% MACs, 32, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(ebn1): BatchNorm2d(32, 0.002% Params, 8.39 MMac, 0.363% MACs, 16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(ebn2): BatchNorm2d(64, 0.004% Params, 4.19 MMac, 0.181% MACs, 32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(ebn3): BatchNorm2d(256, 0.017% Params, 4.19 MMac, 0.181% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(norm3): LayerNorm(320, 0.022% Params, 163.84 KMac, 0.007% MACs, (160,), eps=1e-05, elementwise_affine=True)
(norm4): LayerNorm(512, 0.035% Params, 65.54 KMac, 0.003% MACs, (256,), eps=1e-05, elementwise_affine=True)
(dnorm3): LayerNorm(320, 0.022% Params, 163.84 KMac, 0.007% MACs, (160,), eps=1e-05, elementwise_affine=True)
(dnorm4): LayerNorm(256, 0.017% Params, 524.29 KMac, 0.023% MACs, (128,), eps=1e-05, elementwise_affine=True)
(block1): ModuleList(
(0): shiftedBlock(
53.44 k, 3.631% Params, 54.72 MMac, 2.366% MACs,
(drop_path): Identity(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
(norm2): LayerNorm(320, 0.022% Params, 163.84 KMac, 0.007% MACs, (160,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
53.12 k, 3.609% Params, 54.56 MMac, 2.359% MACs,
(fc1): Linear(25.76 k, 1.750% Params, 26.38 MMac, 1.140% MACs, in_features=160, out_features=160, bias=True)
(dwconv): DWConv(
1.6 k, 0.109% Params, 1.64 MMac, 0.071% MACs,
(dwconv): Conv2d(1.6 k, 0.109% Params, 1.64 MMac, 0.071% MACs, 160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=160)
)
(act): GELU(0, 0.000% Params, 163.84 KMac, 0.007% MACs, approximate='none')
(fc2): Linear(25.76 k, 1.750% Params, 26.38 MMac, 1.140% MACs, in_features=160, out_features=160, bias=True)
(drop): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.0, inplace=False)
)
)
)
(block2): ModuleList(
(0): shiftedBlock(
134.66 k, 9.148% Params, 34.47 MMac, 1.490% MACs,
(drop_path): Identity(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
(norm2): LayerNorm(512, 0.035% Params, 65.54 KMac, 0.003% MACs, (256,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
134.14 k, 9.113% Params, 34.41 MMac, 1.487% MACs,
(fc1): Linear(65.79 k, 4.470% Params, 16.84 MMac, 0.728% MACs, in_features=256, out_features=256, bias=True)
(dwconv): DWConv(
2.56 k, 0.174% Params, 655.36 KMac, 0.028% MACs,
(dwconv): Conv2d(2.56 k, 0.174% Params, 655.36 KMac, 0.028% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=256)
)
(act): GELU(0, 0.000% Params, 65.54 KMac, 0.003% MACs, approximate='none')
(fc2): Linear(65.79 k, 4.470% Params, 16.84 MMac, 0.728% MACs, in_features=256, out_features=256, bias=True)
(drop): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.0, inplace=False)
)
)
)
(dblock1): ModuleList(
(0): shiftedBlock(
53.44 k, 3.631% Params, 54.72 MMac, 2.366% MACs,
(drop_path): Identity(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
(norm2): LayerNorm(320, 0.022% Params, 163.84 KMac, 0.007% MACs, (160,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
53.12 k, 3.609% Params, 54.56 MMac, 2.359% MACs,
(fc1): Linear(25.76 k, 1.750% Params, 26.38 MMac, 1.140% MACs, in_features=160, out_features=160, bias=True)
(dwconv): DWConv(
1.6 k, 0.109% Params, 1.64 MMac, 0.071% MACs,
(dwconv): Conv2d(1.6 k, 0.109% Params, 1.64 MMac, 0.071% MACs, 160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=160)
)
(act): GELU(0, 0.000% Params, 163.84 KMac, 0.007% MACs, approximate='none')
(fc2): Linear(25.76 k, 1.750% Params, 26.38 MMac, 1.140% MACs, in_features=160, out_features=160, bias=True)
(drop): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.0, inplace=False)
)
)
)
(dblock2): ModuleList(
(0): shiftedBlock(
34.56 k, 2.348% Params, 141.56 MMac, 6.120% MACs,
(drop_path): Identity(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
(norm2): LayerNorm(256, 0.017% Params, 524.29 KMac, 0.023% MACs, (128,), eps=1e-05, elementwise_affine=True)
(mlp): shiftmlp(
34.3 k, 2.331% Params, 141.03 MMac, 6.097% MACs,
(fc1): Linear(16.51 k, 1.122% Params, 67.63 MMac, 2.924% MACs, in_features=128, out_features=128, bias=True)
(dwconv): DWConv(
1.28 k, 0.087% Params, 5.24 MMac, 0.227% MACs,
(dwconv): Conv2d(1.28 k, 0.087% Params, 5.24 MMac, 0.227% MACs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128)
)
(act): GELU(0, 0.000% Params, 524.29 KMac, 0.023% MACs, approximate='none')
(fc2): Linear(16.51 k, 1.122% Params, 67.63 MMac, 2.924% MACs, in_features=128, out_features=128, bias=True)
(drop): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.0, inplace=False)
)
)
)
(patch_embed3): OverlapPatchEmbed(
184.8 k, 12.555% Params, 189.07 MMac, 8.174% MACs,
(proj): Conv2d(184.48 k, 12.533% Params, 188.91 MMac, 8.167% MACs, 128, 160, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(norm): LayerNorm(320, 0.022% Params, 163.84 KMac, 0.007% MACs, (160,), eps=1e-05, elementwise_affine=True)
)
(patch_embed4): OverlapPatchEmbed(
369.41 k, 25.097% Params, 94.5 MMac, 4.085% MACs,
(proj): Conv2d(368.9 k, 25.062% Params, 94.44 MMac, 4.083% MACs, 160, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(norm): LayerNorm(512, 0.035% Params, 65.54 KMac, 0.003% MACs, (256,), eps=1e-05, elementwise_affine=True)
)
(decoder1): Conv2d(368.8 k, 25.055% Params, 94.41 MMac, 4.082% MACs, 256, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder2): Conv2d(184.45 k, 12.531% Params, 188.87 MMac, 8.165% MACs, 160, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder3): Conv2d(36.9 k, 2.507% Params, 151.13 MMac, 6.533% MACs, 128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder4): Conv2d(4.62 k, 0.314% Params, 75.76 MMac, 3.275% MACs, 32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(decoder5): Conv2d(2.32 k, 0.158% Params, 152.04 MMac, 6.573% MACs, 16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(dbn1): BatchNorm2d(320, 0.022% Params, 81.92 KMac, 0.004% MACs, 160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(dbn2): BatchNorm2d(256, 0.017% Params, 262.14 KMac, 0.011% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(dbn3): BatchNorm2d(64, 0.004% Params, 262.14 KMac, 0.011% MACs, 32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(dbn4): BatchNorm2d(32, 0.002% Params, 524.29 KMac, 0.023% MACs, 16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(final): Conv2d(34, 0.002% Params, 8.91 MMac, 0.385% MACs, 16, 2, kernel_size=(1, 1), stride=(1, 1))
(soft): Softmax(0, 0.000% Params, 0.0 Mac, 0.000% MACs, dim=1)
)
1.47 M
2.31 GMac
frame_rate 185.77340283438443
默认参数
#mmcv.cnn
mmcv.cnn.get_model_complexity_info(
model,
input_shape, # 必须指定
print_per_layer_stat=True,
as_strings=True,
flush=True,
ost=print # 输出流
)
#ptflops
ptflops.get_model_complexity_info(
model,
input_res, # 可以是tuple或list
print_per_layer_stat=True,
as_strings=True,
flush=False,
ost=print,
verbose=True, # ptflops特有，显示未支持的层
ignore_modules=[...] # ptflops特有，忽略指定层
)
底层逻辑：钩子注入
def my_hook(module, input, output):
# module: 当前层对象
# input: 输入张量的元组
# output: 输出张量
# 示例：统计卷积层的 FLOPs
if isinstance(module, nn.Conv2d):
_, C_in, H, W = input[0].shape
C_out = output.shape[1]
kernel_size = module.kernel_size[0] * module.kernel_size[1]
flops = H * W * C_in * C_out * kernel_size / module.groups
print(f"{module.name}: {flops/1e9:.3f} GFLOPs")
# 注册钩子
handle = model.conv1.register_forward_hook(my_hook)
output = model(input) # 前向时会自动触发钩子
handle.remove() # 移除钩子
浙公网安备 33010602011771号