model.apply(fn)或net.apply(fn)

首先,我们知道pytorch的任何网络net,都是torch.nn.Module的子类,都算是module,也就是模块。
pytorch中的model.apply(fn)会递归地将函数fn应用到父模块的每个子模块submodule,也包括model这个父模块自身。
比如下面的网络例子中。net这个模块有两个子模块,分别为Linear(2,4)Linear(4,8)。函数首先对Linear(2,4)Linear(4,8)两个子模块调用init_weights函数,即print(m)打印Linear(2,4)Linear(4,8)两个子模块。然后再对net模块进行同样的操作。如此完成递归地调用。从而完成model.apply(fn)或者net.apply(fn)
个人水平有限,不足处望指正。
详情可参考
pytorch官网文档.

import torch.nn as nn
@torch.no_grad()
def init_weights(m):
    print(m)
    
net = nn.Sequential(nn.Linear(2,4), nn.Linear(4, 8))
print(net)
print('isinstance torch.nn.Module',isinstance(net,torch.nn.Module))
print(' ')
net.apply(init_weights)

输出

Sequential(
  (0): Linear(in_features=2, out_features=4, bias=True)
  (1): Linear(in_features=4, out_features=8, bias=True)
)
isinstance torch.nn.Module True
Linear(in_features=2, out_features=4, bias=True)
Linear(in_features=4, out_features=8, bias=True)
Sequential(
  (0): Linear(in_features=2, out_features=4, bias=True)
  (1): Linear(in_features=4, out_features=8, bias=True)
)

如果我们想对某些特定的子模块submodule做一些针对性的处理,该怎么做呢。我们可以加入type(m) == nn.Linear:这类判断语句,从而对子模块m进行处理。如下,读者可以细细体会一下。

import torch.nn as nn
@torch.no_grad()
def init_weights(m):
    print(m)
    if type(m) == nn.Linear:
        m.weight.fill_(1.0)
        print(m.weight)
net = nn.Sequential(nn.Linear(2,4), nn.Linear(4, 8))
print(net)
print('isinstance torch.nn.Module',isinstance(net,torch.nn.Module))
print(' ')
net.apply(init_weights)

可以先打印网络整体看看。调用apply函数后,先逐一打印子模块m,然后对子模块进行判断,打印Linear这类子模块m的权重。

Sequential(
  (0): Linear(in_features=2, out_features=4, bias=True)
  (1): Linear(in_features=4, out_features=8, bias=True)
)
isinstance torch.nn.Module True
 
Linear(in_features=2, out_features=4, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.],
        [1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=4, out_features=8, bias=True)
Parameter containing:
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=4, bias=True)
  (1): Linear(in_features=4, out_features=8, bias=True)
)

  • 118
    点赞
  • 192
    收藏
    觉得还不错? 一键收藏
  • 14
    评论
您可以通过定义一个继承自 `tf.keras.Model` 的模型类来创建一个模型实例,并在训练时使用这个模型实例。以下是一个简单的例子: ``` import tensorflow as tf class MyModel(tf.keras.Model): def __init__(self): super(MyModel, self).__init__() self.dense1 = tf.keras.layers.Dense(10, activation='relu') self.dense2 = tf.keras.layers.Dense(1) def call(self, inputs): x = self.dense1(inputs) return self.dense2(x) model = MyModel() ``` 在这个例子中,我们定义了一个继承自 `tf.keras.Model` 的模型类 `MyModel`,并在 `__init__` 方法中定义了两个全连接层。在 `call` 方法中,我们将输入数据传入第一个全连接层,然后将输出传入第二个全连接层。这个模型可以接受任意维度的输入数据,输出一个维度为 1 的标量。 现在,您就可以使用 `model.trainable_variables` 来获取这个模型中的可训练变量了。例如,您可以像下面这样定义一个损失函数和优化器,并使用 `model.trainable_variables` 来更新模型中的变量: ``` optimizer = tf.keras.optimizers.Adam() loss_fn = tf.keras.losses.MeanSquaredError() for x, y in dataset: with tf.GradientTape() as tape: y_pred = model(x) loss = loss_fn(y, y_pred) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) ``` 在这个例子中,我们使用了 `tf.GradientTape` 来计算损失函数的梯度,并使用 `optimizer.apply_gradients` 方法来更新模型中的变量。注意,在计算梯度时,我们使用了 `model.trainable_variables` 来获取模型中的可训练变量。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 14
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值