How to extract features while or after trainnig?

Hi all,
I’m a beginner of pytorch.
I have been struggling with extracting features in forward function.

import torch
import torch.nn as nn

class CrystalGraphConvNet(nn.Module):
    def __init__(self, ...)
        super(CrystalGraphConvNet, self).__init__()
        self.embedding = nn.Linear(orig_atom_fea_len, atom_fea_len)
        self.convs = nn.ModuleList([ConvLayer()])
        self.conv_to_fc = nn.Linear(atom_fea_len, h_fea_len)
        self.conv_to_fc_softplus = nn.Softplus()
        self.fc_out = nn.Linear(h_fea_len, 1)

    def forward(self, ...):
        atom_fea = self.embedding(atom_fea)
        for conv_func in self.convs:
            atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx)
        crys_fea = self.pooling(atom_fea, crystal_atom_idx)
        crys_fea = self.conv_to_fc(self.conv_to_fc_softplus(crys_fea))
        out = self.fc_out(crys_fea)
        return out
from skorch import NeuralNetRegressor
from model import CrystalGraphConvNet

net = NeuralNetRegressor(


I want to extract values of features in the forward function, such as crys_fea.
Please help me out. Thank you!

In plain PyTorch you could simply return crys_fea with out in your forward method.
I’m not sure, if this will work in skorch or if you are limited to a single return value.
Could you try that and see if it’s working?

I set return crys_fea, out in the forward method, but it gives me an error,

TypeError                                 Traceback (most recent call last)
--->         return super(NeuralNetRegressor, self).fit(X, y, **fit_params)
TypeError: mse_loss(): argument 'input' (position 1) must be Tensor, not tuple

That was the issue I was thinking about.
I’ve skimmed through the skorch documentation and it seems some methods like .predict() support multiple outputs, if they are passed as a tuple.
Could you try that?

Hi @ptrblck,

I think .predict() is to prediction using a trained model using .fit(). So I cannot train the model with .predict(). Do you think if there is a way to save this skorch model and load it to extract intermediate features?

Sorry for not being clear enough.
I meant you could try to return both tensors as a tuple in your forward, and see if .fit() will also work similar to .predict().
Let me know if you get stuck.

I think I solved it by subclassing the regressor to make it takes multiple return values in the forward. Thanks!

1 Like