pytorch中apply函数的用法

pytorch中的apply函数是一个高阶函数，可以用来对一个tensor或者一个module中的所有元素进行操作。apply函数的用法如下：

tensor.apply(func)

其中，tensor是要进行操作的tensor，func是一个函数，用来对tensor中的每个元素进行操作。

例如，我们可以使用apply函数对一个tensor中的所有元素进行平方操作：

import torch

x = torch.tensor([1, 2, 3, 4])
x = x.apply(lambda x: x**2)
print(x)

输出结果为：

tensor([ 1, 4, 9, 16])

除了对tensor进行操作，apply函数还可以用来对一个module中的所有参数进行操作。例如，我们可以使用apply函数对一个module中的所有参数进行初始化：

import torch.nn as nn

class MyModule(nn.Module):
def __init__(self):
super(MyModule, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 30)

def init_weights(self):
def init_func(m):
if type(m) == nn.Linear:
nn.init.normal_(m.weight)

self.apply(init_func)

model = MyModule()
model.init_weights()

在上面的代码中，我们定义了一个MyModule类，其中包含两个线性层fc1和fc2。我们使用apply函数对MyModule中的所有线性层的权重进行了初始化。