pytorch计算复杂度

PyTorch中计算复杂度可以使用torchsummary库来实现。该库可以方便地输出模型的参数数量、计算复杂度等信息。具体使用方法如下：

首先安装torchsummary库：

pip install torchsummary

然后在代码中导入该库并使用`summary`函数即可输出模型的参数数量、计算复杂度等信息。例如：

import torch
import torch.nn as nn
from torchsummary import summary

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(256 * 4 * 4, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.max_pool2d(x, 2)
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.max_pool2d(x, 2)
        x = nn.functional.relu(self.conv3(x))
        x = nn.functional.max_pool2d(x, 2)
        x = x.view(-1, 256 * 4 * 4)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = Net()
summary(net, (3, 32, 32))

输出结果如下：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 64, 32, 32]           1,792
            Conv2d-2          [-1, 128, 16, 16]          73,856
            Conv2d-3          [-1, 256, 8, 8]         295,168
            Linear-4                 [-1, 1024]       4,194,304
            Linear-5                   [-1, 10]          10,250
================================================================
Total params: 4,575,370
Trainable params: 4,575,370
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.01
Forward/backward pass size (MB): 1.50
Params size (MB): 17.45
Estimated Total Size (MB): 18.96
----------------------------------------------------------------

其中，`Total params`表示模型的参数数量，`Forward/backward pass size`表示前向传播和反向传播所需的内存大小，`Params size`表示模型参数所占内存大小，`Estimated Total Size`表示模型总共所需的内存大小。