Search code examples
pythonneural-networkpytorchskorch

How do I retrieve the weights estimated from a neural net using skorch?


I've trained a simple neural net using skorch to make it sklearn compatible and I would like to know how to retrieve the actual estimated weights.

Here's a replicable example of what I need.

The neural net presented here uses 10 features, has one hidden layer of 2 nodes, uses ReLu activation functions and linearly combines the output of the 2 nodes.

import torch
import numpy as np
from torch.autograd import Variable

# Create example data 
np.random.seed(2022)
train_size = 1000
n_features= 10

X_train = np.random.rand(n_features, train_size).astype("float32")

l2_params_1 = np.random.rand(1,n_features).astype("float32")
l2_params_2 = np.random.rand(1,n_features).astype("float32")

l1_X = np.matmul(l2_params_1, X_train)
l2_X = np.matmul(l2_params_2, X_train)

y_train = l1_X + l2_X

# Defining my NN
class NNModule(torch.nn.Module):
    def __init__(self, in_features):
        super(NNModule, self).__init__()
        self.l1 =  torch.nn.Linear(in_features, 2)
        self.a1 = torch.nn.ReLU()
        self.l2 =  torch.nn.Linear(2, 1)

    def forward(self, x):
        x = self.l1(x)
        x = self.a1(x)

        return self.l2(x)

# Initialize the NN
torch.manual_seed(200)

model = NNModule(in_features = 10)

model.l1.weight.data.uniform_(0.0, 1.0)
model.l1.bias.data.uniform_(0.0, 1.0)

# Define criterion and optimizer
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Train the NN
torch.manual_seed(200)

for epoch in range(100):
    inputs = Variable(torch.from_numpy(np.transpose(X_train)))
    labels = Variable(torch.from_numpy(np.transpose(y_train)))

    optimizer.zero_grad()

    outputs = model(inputs)

    loss = criterion(outputs, labels)

    loss.backward()

    optimizer.step()

The parameters at which I'm arriving are the following:

list(model.parameters())

[Output]: 
[Parameter containing:
 tensor([[0.8997, 0.8345, 0.8284, 0.6950, 0.5949, 0.1217, 0.9067, 0.1824, 0.8272,
          0.2372],
         [0.7525, 0.6577, 0.4358, 0.6109, 0.8817, 0.5429, 0.5263, 0.7531, 0.1552,
          0.7066]], requires_grad=True),
 Parameter containing:
 tensor([0.6617, 0.1079], requires_grad=True),
 Parameter containing:
 tensor([[0.9225, 0.8339]], requires_grad=True),
 Parameter containing:
 tensor([0.0786], requires_grad=True)]

Now, to wrap my NNModule with skorch, I'm using this:

from skorch import NeuralNetRegressor

torch.manual_seed(200)
net = NeuralNetRegressor(
    module=NNModule(in_features=10),
    criterion=torch.nn.MSELoss,
    optimizer=torch.optim.SGD,
    optimizer__lr=0.01,
    max_epochs=100,
    verbose=0
)

net.fit(np.transpose(X_train), np.transpose(y_train))

And I'd like to retrieve the weights obtained in the training. I've used dir(net) to see if the weights are stored in any attributes to no avail.


Solution

  • To retrieve the weights one needs to output them like this:

    list(net.module.parameters())