在复杂模型上使用Pytorch进行修剪

时间:2020-07-20 18:20:25

标签: python machine-learning pytorch google-colaboratory pruning

所以我正在尝试使用torch.nn.utils.prune.global_unstructured

我在一个简单的模型上做到了,而且行得通。 model.cov2或其他图层即可。我正在尝试使用(嵌套的)模型吗?我收到以下错误消息:

AttributeError: 'CNN' object has no attribute 'conv1'

和其他错误。我尝试了所有方法来访问此深层cov1,但是无法。

您可以在下面找到型号代码:

class CNN(nn.Module):
def __init__(self):
    """CNN Builder."""
    super(CNN, self).__init__()

    self.conv_layer = nn.Sequential(

        # Conv Layer block 1
        nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1),
        nn.BatchNorm2d(32),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2, stride=2),

        # Conv Layer block 2
        nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
        nn.BatchNorm2d(128),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2, stride=2),
        nn.Dropout2d(p=0.05),

        # Conv Layer block 3
        nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1),
        nn.BatchNorm2d(256),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2, stride=2),
    )


    self.fc_layer = nn.Sequential(
        nn.Dropout(p=0.1),
        nn.Linear(4096, 1024),
        nn.ReLU(inplace=True),
        nn.Linear(1024, 512),
        nn.ReLU(inplace=True),
        nn.Dropout(p=0.1),
        nn.Linear(512, 100)
    )


def forward(self, x):
    """Perform forward."""
    # conv layers
    x = self.conv_layer(x)
    # flatten
    x = x.view(x.size(0), -1)
    # fc layer
    x = self.fc_layer(x)
    return x

如何在此模型上应用修剪?

2 个答案:

答案 0 :(得分:0)

您的模块不是名称“conv1”或“conv2”,您可以使用 named_modules 生成器查看名称。从上面,你有一个“conv_stem”,它可以被索引为 model.conv_stem[0] 来访问。您可以遍历模块来创建一个字典,如:

parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
    (model.fc2, 'weight'),
    (model.fc3, 'weight'), )

并将其传入。查看更多信息:https://colab.research.google.com/github/pytorch/tutorials/blob/gh-pages/_downloads/f40ae04715cdb214ecba048c12f8dddf/pruning_tutorial.ipynb#scrollTo=UVFjM079F0Oi

答案 1 :(得分:0)

使用此方法查看图层名称

for layer_name, param in model.named_parameters():
    print(f"layer name: {layer_name} has {param.shape}")

并将这些名称传递给 prune 方法

例如,在prune.random_unstructured(module_name, name="weight", amount=0.3)