Pytorch获取模型所有层

最新推荐文章于 2024-07-30 19:10:49 发布

非典型废言

最新推荐文章于 2024-07-30 19:10:49 发布

阅读量3.9k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： TensorFlow

本文链接：https://blog.csdn.net/sinat_35821976/article/details/103733007

TensorFlow 专栏收录该内容

10 篇文章

订阅专栏

本文深入探讨了如何使用Python和PyTorch库获取深度学习模型的每一层及其模块，通过递归方式展开模型的层级结构，帮助读者理解模型内部组成。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

def getLayers(model):
    """
    get each layer's name and its module
    :param model:
    :return: each layer's name and its module
    """
    layers = []

    def unfoldLayer(model):
        """
        unfold each layer
        :param model: the given model or a single layer
        :param root: root name
        :return:
        """

        # get all layers of the model
        layer_list = list(model.named_children())
        for item in layer_list:
            module = item[1]
            sublayer = list(module.named_children())
            sublayer_num = len(sublayer)

            # if current layer contains sublayers, add current layer name on its sublayers
            if sublayer_num == 0:
                layers.append(module)
            # if current layer contains sublayers, unfold them
            elif isinstance(module, torch.nn.Module):
                unfoldLayer(module)

    unfoldLayer(model)
    return layers