CODEWITHSUNDEEP
pythondeep-learningneural-networkpytorch
Penguin
Penguin

Reputation: 2411

Efficient way to get the weights of a PyTorch NN model as a tensor

I have a simple 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()

Is there a more efficient way to get the weights of this network (while keeping the gradients) than iterate through every single one like this

for w_arr in net.parameters():

or

list(net.parameters())

Since the latter doesn't maintain the gradients (it converts it into a list)

Upvotes: 1

Views: 1575

Answers (1)

Ivan
Ivan

Reputation: 40628

You can use the torch.nn.utils.parameters_to_vector utility function.

>>> 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>)

It will convert a parameter generator into a flat tensor while retaining gradients, which corresponds to a concatenation of all parameter tensors flattened.

Upvotes: 2

Related Questions