I’m using torch to write some non neural net code, and I have something like:
class Worker: def __init__(self, fixed_matrix, fixed_neural_net): self.fixed_matrix = fixed_matrix self.fixed_neural_net = fixed_neural_net def do_work(self, data): # something with self.fixed_matrix,fixed_neural_net
I’d like to be able to use
worker.to("cuda:1"), and things like this, and one approach would be to subclass
nn.Module and use
self.register_buffer(...) to store
fixed_neural_net is already a
Module instance here (a neural net which is not training, and has requires_grad=False).
However, I don’t want all of the extra bells and whistles that
Module brings with it related to NN training: things like
forward(), the parameters, the forward and backward hooks, etc. To be clear: gradients will not be used at all in this setting.
I thought this kind of case might be common, so I was wondering if there is a simple way to just take the parts of
Module which I want (really, all of the logic related to
_modules) and forget the rest. Or, maybe the rest is no big deal to have around? The main thing I don’t want to do is have to call
requires_grad_(False) every time I initialize this Worker, I guess.