如何在tensorflow中使用pytorch的ModuleList、ModuleDict功能

今夜月色真好

已于 2023-03-24 15:25:37 修改

阅读量463

点赞数

分类专栏： github闲逛录文章标签： pytorch tensorflow 深度学习 Powered by 金山文档

于 2023-03-24 15:21:48 首次发布

本文链接：https://blog.csdn.net/h8215066/article/details/129750977

版权

这篇博客介绍了如何在TensorFlow中模仿PyTorch的nn.ModuleList和nn.ModuleDict，这两个容器在构建神经网络时提供了灵活性。尽管TensorFlow官方未提供这些工具，但有第三方库实现了类似功能，允许与TensorFlow代码无缝结合。文章通过代码示例展示了使用方法，并保证了与原生TensorFlow的兼容性。

摘要由CSDN通过智能技术生成

在使用pytorch时构建网络时，我们经常会使用到nn.ModuleList和nn.ModuleDict来帮助我们保存要参与构建的神经网络层，这两类被称之为Module 容器（containers）。

nn.ModuleList

nn.ModuleList 从名字我们可以看出，是可以将神经网络层（比如 nn.Conv2d, nn.Linear 之类的）保存在类似于List一样的容器中，然后可以使用List同名的方法来进行操作，比如index取值、append、extend等方法。但与原始的List不同的点在于nn.ModuleList中的Module会注册到网络，并自动收集Module中的可训练parameters。在保存的时候，并不直接指定Modules之间的连接关系，而在后续的前向计算中另行指定。

下面我们看一段代码：

import torch
from torch import nn
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10,1), nn.Linear(20,10), nn.Linear(30,20)])
    def forward(self, x):
        x = self.linears[2](x)
        x = self.linears[1](x)
        x = self.linears[0](x) 
        return x

net = MLP()
pri