How to create a network with multiple forward routines

1

I have a bit of a special situation that I hope someone can help with.

I wish to create a neural network y=f(x), where f is the network, x is input and y is output. However the network should be made in such a way that it can effectively be split in two parts.
lets say f() = g(h()). The idea is that h is the feature extracting part of the network, while g is the classifying part of the network.

I now wish to be able to call either g, h, or f on my data. However I can’t quite figure out how to do this, since each neural network seems to be able to only have 1 forward routine.

I guess one way to achieve this would be to write two separate neural networks, g and h. I take it the autograd wouldn’t have any problem with several neural networks like that. Would there be any performance hits in this approach or is there another more elegant approach?

1 Like
2

Hi,

I would do the following:

class H:
  def __init__(self, *args):
    # init H

  def forward(self, inp):
    # forward of H
    return out

class G:
  def __init__(self, *args):
    # init G

  def forward(self, inp):
    # forward of G
    return out

class F:
  def __init__(self, *args):
    # init F
    self.g = G()
    self.h = H()

  def forward(self, inp):
    # forward of F
    return self.g(self.h(inp))

model = F()

# forward the whole model
model(inp)

# foward only g
model.g(inp)
# get only g's parameters
model.g.paremeters()

# foward only h
model.h(inp)
4 Likes
3

I will try to do something similar to this, unfortunately I don’t think it will work quite as elegantly, since I want to offload the model to cuda in nn.dataparallel, at which point I don’t believe I can access any of the subparameters or functions of the model.