How to make custom method in nn.Module work with GPUs

moskomule: <p>I’m trying to implement simple res-net like below and it works with CPU.</p>
<pre><code class="lang-auto">class ResNet(nn.module):
      def __init__(self):
        super(Net, self).__init__()
        self.conv_1 = nn.Conv2d(3, 64, 4, stride=2)
        self.bn_1 = nn.BatchNorm2d(64)
        self.res_1 = self.__res_block(64, [32,32,128], True)
         ...
    def forward(self, x):
        x = self.conv_1(x)
        x = F.relu(self.bn_1(x))
        x = F.max_pool2d(x, 2, 2)
        x = self.res_1(x)
        ...

    def __res_block(self, in_channels,
                    nb_filters, right=False):
            def __res_base(_in_channels, out_channels,
                       kernel_size=1, padding=0):
            def g(x):
                x = nn.Conv2d(in_channels=_in_channels,
                              out_channels=out_channels,
                              kernel_size=kernel_size,
                              padding=padding)(x)
                x = nn.BatchNorm2d(num_features=out_channels)(x)
                return x
            return g

        def f(x):
            y = F.relu(__res_base(in_channels, nb_filters[0])(x))
            y = F.relu(__res_base(nb_filters[0], nb_filters[1],
                                  kernel_size=3, padding=1)(y))
            y = F.relu(__res_base(nb_filters[1], nb_filters[2])(y))
            if right is True:
                x = __res_base(in_channels, nb_filters[2])(x)
            return F.relu(x+y)
        return f
</code></pre>
<p>but it doesn’t work with GPU and throws TypeError,</p>
<pre><code class="lang-auto">TypeError: _cudnn_convolution_full_forward received an invalid combination of arguments 
- got (torch.cuda.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.cuda.FloatTensor, tuple, tuple, int, bool), 
but expected (torch.cuda.RealTensor input, torch.cuda.RealTensor weight, torch.cuda.RealTensor bias, 
torch.cuda.RealTensor output, std::vector&lt;int&gt; pad, std::vector&lt;int&gt; stride, int groups, bool benchmark)
</code></pre>
<p>I think this is because the parameters in <code>__res_block</code> doesn’t pass to GPU but how can I pass them to GPU? Thank you.</p>

albanD: <p>Hi,</p>
<p>From the error message it looks like you give him <code>torch.FloatTensor</code> as input.<br>
If your model is on GPU, you should give him <code>torch.cuda.FloatTensor</code> (EDIT: sorry wrong name for cuda tensors).<br>
You can convert your inputs to cuda by doing: <code>input = input.cuda()</code> before forwarding it through the network.</p>

moskomule: <p>I uploaded the code <a href="https://github.com/moskomule/pytorch_learning/blob/master/src/resnet.py" rel="nofollow noopener">here</a>.</p>
<p>The code except the module above is same as my other examples and they work with GPUs. So I think the cause is in the <code>__res_block</code>.</p>
<p>And, does <code>torch.CudaFloatTensor</code>  equal to <code>torch.cuda.RealTensor</code>?</p>

alexis-jacq: <p>In your <code>__res_block</code> you create new objects that are not composed by <code>torch.CudaFloatTensors</code></p>
<p>Did you try just in <code>def g(x)</code>:<br>
<code>return x.cuda()</code> instead of <code>return x</code></p>

albanD: <p>Ho I see,</p>
<p>I think you want both your <code>__res_block</code> and <code>__res_base</code> functions to return an instance of <code>Module</code>s.<br>
Otherwise their parameters won’t be recognized as being part of the main network.<br>
You want both of them to be classes that subclass the <code>nn.Module</code> class.<br>
In each, you want in the init to initialize the operations and store them in <code>self</code> as done <a href="https://github.com/moskomule/pytorch_learning/blob/659a26577c64ce6083d24a7cf25309cfc6e6c8e1/src/resnet.py#L33-L34">here</a> then a forward pass that just use them as done <a href="https://github.com/moskomule/pytorch_learning/blob/659a26577c64ce6083d24a7cf25309cfc6e6c8e1/src/resnet.py#L48-L49">here</a>.</p>

apaszke: <p>There’s no such thing as <code>torch.CudaFloatTensor</code>, only <code>torch.cuda.FloatTensor</code>.</p>
<p>Also, as <a class="mention" href="/u/alband">@albanD</a> said, you can’t return a closure from the <code>__res_block</code> and expect the modules inside it to be recognized as part of the model. Just create a new <code>nn.Module</code> subclass. You can see how ResNets are implemented in <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py">torchvision</a>.</p>

moskomule: <p><a class="mention" href="/u/alexis-jacq">@alexis-jacq</a><br>
thank you for your advice, but that didn’t work  because the input to <code>g(x)</code> is already <code>torch.cuda.**Tensor</code> and the problem is the weights etc are not <code>torch.cuda.**Tensor</code>.</p>
<p><a class="mention" href="/u/alband">@albanD</a> <a class="mention" href="/u/apaszke">@apaszke</a><br>
thank you two, ok I’ll create blocks as <code>nn.Module</code> subclasses like the tochvision implement.</p>
<p><a class="mention" href="/u/apaszke">@apaszke</a><br>
Is <code>torch.cuda.RealTensor</code> in the TypeError message equivalent to <code>torch.cuda.FloatTensor</code>? I couldn’t find <code>RealTensor</code> in the doc.</p>

smth: <p><a class="mention" href="/u/moskomule">@moskomule</a> <code>torch.cuda.RealTensor</code> refers to any <code>torch.cuda.*Tensor</code>, and in this context, is probably a <code>torch.cuda.FloatTensor</code></p>

moskomule: <p>thank you all, I updated <a href="https://github.com/moskomule/pytorch_learning/blob/master/src/resnet.py" rel="nofollow noopener">the code</a> and it works well.</p>
<p>So briefly, <em>the instances of subclasses of <code>nn.Module</code> are able to be backproped</em>.</p>

apaszke: <p>No, <code>nn.Module</code> subclasses only add support for some convenience methods like <code>.parameters()</code>, <code>.cuda()</code> and some others. It’s possible to implement neural networks in autograd without modules, and some people actually prefer this style. But this means that you’ll need some additional helpers and custom data structures for handling parameters.</p>

marioviti: <p>All links are dead<br>
could you post here the solution?</p>