torch.nn.Module.apply
是 PyTorch 中用于递归地应用函数到模型的所有子模块的方法。它允许对模型中的每个子模块进行操作,比如初始化权重、改变参数类型等。
以下是关于 torch.nn.Module.apply 的示例:
Module.apply(fn)
import torch
import torch.nn as nn
# 自定义初始化函数
def init_weights(module):
if isinstance(module, nn.Conv2d):
nn.init.xavier_uniform_(module.weight)
elif isinstance(module, nn.Linear):
nn.init.normal_(module.weight, mean=0, std=0.01)
nn.init.constant_(module.bias, 0)
# 定义一个神经网络模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv = nn.Conv2d(3, 16, 3)
self.fc = nn.Linear(16 * 28 * 28, 10)
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
# 创建模型实例
model = MyModel()
# 对模型的所有子模块应用初始化权重的函数
model.apply(init_weights)
import torch
import torch.nn as nn
# 自定义函数:将所有参数类型转换为 float 类型
def convert_to_float(module):
if hasattr(module, 'weight'):
module.weight = nn.Parameter(module.weight.float())
if hasattr(module, 'bias'):
module.bias = nn.Parameter(module.bias.float())
# 创建一个预训练的模型
pretrained_model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True)
# 将预训练模型的参数类型转换为 float
pretrained_model.apply(convert_to_float)
torch.nn.Module.apply
提供了一种方便的方式,允许对模型的每个子模块应用自定义函数,从而进行各种操作,如初始化权重、参数类型转换等。
注意事项:
torch.nn.Module.apply
方法是一个强大的工具,允许你对模型的每个子模块进行操作,从而实现初始化、类型转换、参数修改等一系列功能。通过传入不同的操作函数,你可以灵活地定制和修改模型。