pytorch笔记：ModuleDict

最新推荐文章于 2024-05-31 11:56:36 发布

原创

最新推荐文章于 2024-05-31 11:56:36 发布 · 579 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #笔记 #人工智能

本文介绍了PyTorch中的nn.ModuleDict，一个用于存储和动态管理nn.Module对象的容器，它提供了组织模块、自动参数注册以及常见的操作方法，如添加、删除和访问模块。通过实例展示了如何在模型中使用ModuleDict来简化代码结构。

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

1 介绍

在 PyTorch 中，nn.ModuleDict 是一个方便的容器，用于存储一组子模块（即 nn.Module 对象）的字典
这个容器主要用于动态地管理多个模块，并通过键来访问它们，类似于 Python 的字典

2 特点

组织性
- nn.ModuleDict 提供了一种将多个模块有序组织在一起的方法。
- 这有助于让代码更加结构化，易于理解和维护
动态操作
- 可以像操作普通字典那样添加或删除模块
- 例如使用 module_dict['key'] = module 添加模块，使用 del module_dict['key'] 删除模块
自动参数注册
- 当将模块添加到 ModuleDict 中时，它们的参数会自动注册到整个网络中，确保在模型训练时这些参数可以被识别和更新

3 例子

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layers = nn.ModuleDict({
            'linear': nn.Linear(10, 20),
            'activation': nn.R