AdamCobb / hamiltorch

PyTorch-based library for Riemannian Manifold Hamiltonian Monte Carlo (RMHMC) and inference in Bayesian neural networks
BSD 2-Clause "Simplified" License
426 stars 63 forks source link

Conversion of model to functional breaks with internal methods #8

Open chrhenning opened 4 years ago

chrhenning commented 4 years ago

Hey,

thanks so much for providing the repo. At the moment, it seems that the code can't deal with PyTorch models that call internal methods inside the forward method. I think it would be a great feature if such models could be supported in the future.

Here is an example model, that wouldn't work (assuming that I use the code correctly):

class SimpleMLP(torch.nn.Module):
    def __init__(self, n_in=1, n_out=1, layers=()):
        super(SimpleMLP, self).__init__()

        ### Hidden layers.
        prev_dim = n_in
        self._hidden = torch.nn.Sequential()

        for i in range(len(layers)):
            self._hidden.add_module('hidden_layer_%d' % i,
                torch.nn.Linear(prev_dim, layers[i]))
            self._hidden.add_module('relu_%d' % i, torch.nn.ReLU())
            prev_dim = layers[i]

        ### Generate linear output layer.
        h_size = prev_dim
        self._out_layer = torch.nn.Linear(h_size, n_out)

    def forward(self, x):
        return self._int_forward(x)

    def _int_forward(self, x):
        h = self._hidden(x)
        return self._out_layer(h)

The reference to the method _int_forward is not considered in the function hamiltorch.util.make_functional.

Again, thanks for the great repo, Christian

AdamCobb commented 4 years ago

Hi Christian,

Thanks very much for bringing this example to my attention! I will make sure that I look into that soon.

I plan on making a few code updates soon and will also try to test out your example when I push the next updates.

Thanks very much for using the repo!

All the best,

Adam