Блог

У меня есть простой NN:

 import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(1, 5)
        self.fc2 = nn.Linear(5, 10)
        self.fc3 = nn.Linear(10, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)        
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Model()
 

Есть ли более эффективный способ получить веса этой сети (при сохранении градиентов), чем перебирать каждый из них подобным образом

 for w_arr in net.parameters():
 

или

 list(net.parameters())
 

Поскольку последний не поддерживает градиенты (он преобразует его в список)

Вы можете использовать torch.nn.utils.parameters_to_vector служебную функцию.

 >>> net(torch.rand(1, 1, requires_grad=True)).mean().backward()

>>> from torch.nn.utils import parameters_to_vector
>>> parameters_to_vector(net.parameters())
tensor([-0.8196, -0.7785, -0.2459,  0.4670, -0.9747,  0.1994,  0.7510, -0.6452,
         0.4948,  0.3376,  0.2641, -0.0707,  0.1282, -0.2944,  0.1337,  0.0461,
        -0.1491,  0.2985,  0.3031,  0.3566,  0.0058,  0.0157, -0.0712,  0.3874,
         0.2870, -0.3829,  0.1178, -0.3901, -0.0425, -0.1603,  0.0408,  0.3513,
         0.0289, -0.3374, -0.1820,  0.3684, -0.3069,  0.0312, -0.4205,  0.1456,
         0.2833,  0.0589, -0.2229, -0.1753, -0.1829,  0.1529,  0.1097,  0.0067,
        -0.2694, -0.2176,  0.2292,  0.0529, -0.2617,  0.0736,  0.1617,  0.0438,
         0.2387,  0.3278, -0.0536, -0.2875, -0.0869,  0.0770, -0.0774, -0.1909,
         0.2803, -0.3237, -0.3851, -0.2241,  0.2838,  0.2202,  0.3057,  0.0128,
        -0.2650,  0.1660, -0.2961, -0.0123, -0.2106, -0.1021,  0.1135, -0.1051,
         0.1735], grad_fn=<CatBackward>)
 

Он преобразует генератор параметров в плоский тензор с сохранением градиентов, что соответствует конкатенации всех сглаженных тензоров параметров.