PyTorch 获取模型的所有层

Question

在没有任何nn.Sequence 分组的情况下，获取 pytorch 模型并获取所有层列表的最简单方法是什么？例如，有更好的方法吗？

import pretrainedmodels

def unwrap_model(model):
    for i in children(model):
        if isinstance(i, nn.Sequential): unwrap_model(i)
        else: l.append(i)

model = pretrainedmodels.__dict__['xception'](num_classes=1000, pretrained='imagenet')
l = []
unwrap_model(model)            
            
print(l)
    

Answer 1

您可以使用modules() 方法遍历模型的所有模块（包括每个Sequential 中的模块）。这是一个简单的例子：

>>> model = nn.Sequential(nn.Linear(2, 2), 
                          nn.ReLU(),
                          nn.Sequential(nn.Linear(2, 1),
                          nn.Sigmoid()))

>>> l = [module for module in model.modules() if not isinstance(module, nn.Sequential)]

>>> l

[Linear(in_features=2, out_features=2, bias=True),
 ReLU(),
 Linear(in_features=2, out_features=1, bias=True),
 Sigmoid()]

Answer 2

我将它用于更深层次的模型，并非所有块都来自 nn.sequential。

def get_children(model: torch.nn.Module):
    # get children form model!
    children = list(model.children())
    flatt_children = []
    if children == []:
        # if model has no children; model is last child! :O
        return model
    else:
       # look for children from children... to the last child!
       for child in children:
            try:
                flatt_children.extend(get_children(child))
            except TypeError:
                flatt_children.append(get_children(child))
    return flatt_children

Answer 3

我是这样做的：

def flatten(el):
    flattened = [flatten(children) for children in el.children()]
    res = [el]
    for c in flattened:
        res += c
    return res

cnn = nn.Sequential(Custom_block_1, Custom_block_2)
layers = flatten(cnn)

Answer 4

如果您想要命名dict 中的图层，这是最简单的方法：

named_layers = dict(model.named_modules())

这会返回类似：

{
    'conv1': <some conv layer>,
    'fc1': < some fc layer>,
     ### and other layers 
}

例子：

import torchvision.models as models

model = models.inception_v3(pretrained = True)
named_layers = dict(model.named_modules())

Answer 5

如果你想要一个以名称为键、模块为值的嵌套字典，例如：

{'conv1': Conv2d(...),
 'bn1': BatchNorm2d(...),
 'block1':{
    'group1':{
        'conv1': Conv2d(...),
        'bn1': BatchNorm2d(...),
        'conv2': Conv2d(...),
        'bn2': BatchNorm2d(...),
    },
    'group2':{ ...
    }, ...
}

您可以结合 Kees 和 Mayukh Deb 的答案得到：

def nested_children(m: torch.nn.Module):
    children = dict(m.named_children())
    output = {}
    if children == {}:
        # if module has no children; m is last child! :O
        return m
    else:
        # look for children from children... to the last child!
        for name, child in children.items():
            try:
                output[name] = nested_children(child)
            except TypeError:
                output[name] = nested_children(child)
    return output

Answer 6

这是我的方法，你一般可以在这里输入任何模型，它会返回一个包含所有torch.nn.*事物的列表

def flatten_model(modules):
    def flatten_list(_2d_list):
        flat_list = []
        # Iterate through the outer list
        for element in _2d_list:
            if type(element) is list:
                # If the element is of type list, iterate through the sublist
                for item in element:
                    flat_list.append(item)
            else:
                flat_list.append(element)
        return flat_list

    ret = []
    try:
        for _, n in modules:
            ret.append(loopthrough(n))
    except:
        try:
            if str(modules._modules.items()) == "odict_items([])":
                ret.append(modules)
            else:
                for _, n in modules._modules.items():
                    ret.append(loopthrough(n))
        except:
            ret.append(modules)
    return flatten_list(ret)