便携使用pytorch进行剪枝的代码

您好，以下是使用 PyTorch 进行剪枝的示例代码：

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

# 剪枝卷积层的权重
prune.random_unstructured(net.conv1, name="weight", amount=.2)
prune.random_unstructured(net.conv2, name="weight", amount=.2)

# 剪枝全连接层的权重
prune.l1_unstructured(net.fc1, name="weight", amount=.2)
prune.l1_unstructured(net.fc2, name="weight", amount=.2)

# 剪枝后，需要移除被剪枝的权重
prune.remove(net.conv1, "weight")
prune.remove(net.conv2, "weight")
prune.remove(net.fc1, "weight")
prune.remove(net.fc2, "weight")

希望能对您有所帮助。