Unable to store the U-net model in cuda device

username: Alan_lam
cooked: <p>I am building a CNN model with residual network achidecture for input images with different dimension. At the same time, I want to make sure the output has the same dimension as the input.</p>
<p>In the forward pass of the Generator object, I declare a list object <code>size_list</code> to store the dimension of the tensors during its down-sampling steps. So that when the tensor is unsampling, I can pass the required size from the list to the forward function of ConvTranspose2D in the argument <code>output_size</code>.</p>
<p>However, once I added this list object, I run into this error even though I explicitly pass both input and model into my cuda device.</p>
<p><code>RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same</code></p>
<p>The error smply would not show up if I don’t use the list object, or run the model in cpu device. I don’t know how the list object influence my model getting into the cuda device. Any help is appreciated. Thank you.</p>
<pre><code class="lang-auto">class Transpose_Block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, output_size, padding=1):
        super(Transpose_Block, self).__init__()
        self.output_size = output_size
        self.block = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)

    def centercrop(self, x):
        dimension = x.shape
        top = 0 + math.floor((dimension[-2] - self.output_size[0]) / 2)
        bottom = dimension[-2] - math.ceil((dimension[-2] - self.output_size[0]) / 2)
        left = 0 + math.floor((dimension[-1] - self.output_size[1]) / 2)
        right = dimension[-1] - math.ceil((dimension[-1] - self.output_size[1]) / 2)
        print(top, bottom, left, right)
        return x[:,:, top:bottom, left:right]

    def forward(self, x):
        try:
            out = self.block(x, output_size=self.output_size)
        except ValueError:
            out = self.block(x)
            out = self.centercrop(out)
            print('manual crop')
        return out


class Generator(nn.Module):
    def __init__(self, input_channels=INPUT_CHANNELS, channels_list=[64, 128, 256, 512], n_res=NUM_RES):
        super(Generator, self).__init__()
        self.input_channels = input_channels
        self.channels_list = channels_list
        self.list = []

        self.initial = nn.Sequential(
            nn.Conv2d(input_channels, channels_list[0], kernel_size=7, stride=1, padding=3, padding_mode='reflect'),
            nn.InstanceNorm2d(channels_list[0]),
            nn.ReLU(),
        )

        self.down_layers = nn.ModuleList()
        self.res_layers = nn.ModuleList()

        self.down_layers.append(self.initial)
        for down in range(len(channels_list)-1):
            self.down_layers.append(Gen_Block(channels_list[down], channels_list[down + 1]))
        for res in range(n_res):
            self.res_layers.append(Res_Block(channels_list[-1]))


    def forward(self, x):
        # size_tensor = torch.empty((len(self.channels_list),2))
        for (index, down) in enumerate(self.down_layers):
            # size_tensor[index] = torch.tensor(x.size()[-2:])
            self.list.append(x.size()[-2:])
            x = down(x)
        for res in self.res_layers:
            x = res(x)
        for up in range(len(self.channels_list)-1, 0, -1):
            x = Transpose_Block(self.channels_list[up], self.channels_list[up-1], kernel_size=3, stride=2, output_size=self.list[up], padding=1)(x)
            x = nn.InstanceNorm2d(self.channels_list[up-1])(x)
            x = nn.ReLU()(x)
        x = Transpose_Block(self.channels_list[0], self.input_channels, kernel_size=7, stride=1, output_size=self.list[0], padding=3)(x)
        x = nn.Tanh()(x)

        return x


def gen_test():
    x = torch.randn(1, 1, 100, 150).to(DEVICE)
    model = Generator().to(DEVICE)
    print(model)
    print(model(x).shape)



if __name__ == '__main__':
    gen_test()
    pass
</code></pre>

username: KFrank
cooked: <p>Hi Alan!</p>
<aside class="quote no-group quote-modified" data-username="Alan_lam" data-post="1" data-topic="145623" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alan_lam/48/43715_2.png" class="avatar"> Alan_lam:</div>
<blockquote>
<pre><code class="lang-auto">class Transpose_Block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, output_size, padding=1):
        super(Transpose_Block, self).__init__()
        self.output_size = output_size
        self.block = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
...

class Generator(nn.Module):
...
    def forward(self, x):
...
        for up in range(len(self.channels_list)-1, 0, -1):
            x = Transpose_Block(self.channels_list[up], self.channels_list[up-1], kernel_size=3, stride=2, output_size=self.list[up], padding=1)(x)
            x = nn.InstanceNorm2d(self.channels_list[up-1])(x)
            x = nn.ReLU()(x)
        x = Transpose_Block(self.channels_list[0], self.input_channels, kernel_size=7, stride=1, output_size=self.list[0], padding=3)(x)
        x = nn.Tanh()(x)

        return x


def gen_test():
    x = torch.randn(1, 1, 100, 150).to(DEVICE)
    model = Generator().to(DEVICE)
...
</code></pre>
</blockquote>
</aside>
<p>When you move your <code>model</code> (a <code>Generator</code>) to your gpu, pytorch goes<br>
through (at least roughly speaking) your <code>model</code>’s properties, and moves<br>
the various layers, etc., to the gpu.</p>
<p>The problem is that when you do this no instance of <code>Transpose_Block</code><br>
has yet been created, so there is nothing for pytorch to move to the gpu<br>
nor to even know about.</p>
<p>Instead, in <code>Generator</code>’s <code>forward()</code> method, you create <code>Transpose_Block</code><br>
on the fly, so it – and its “layer,” <code>ConvTranspose2d</code>, is created (by default)<br>
on the cpu, and never moved to the gpu.  Then when the gpu tensors try to<br>
flow through it, you get your error.</p>
<p>As an aside, every time you create your <code>Transpose_Block</code> on the fly,<br>
you instantiate a new <code>ConvTranspose2d</code> with new, different, randomly<br>
initialized values for its <code>weight</code> and <code>bias</code>.  (These parameters are never<br>
added to your optimizer, so they never get updated by the optimizer – but<br>
that wouldn’t matter because you discard them anyway.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: Alan_lam
cooked: <p>Thank you so much for your reminder!! I didn’t think about this.</p>
<p>I guess the <code>Transpose_Block</code> cannot be instantiatized inside forward method too because it otherwise would be re-created in every forward pass? In such case where I can only get the <code>size_list</code> after passing the forward method, what is the solution if I still want to pass the <code>output_size</code> argument when instantiatizing the <code>Transpose_Block</code>? Do I have to manually calculate the <code>size_list</code> according to the <code>kernel_size</code>, <code>padding</code> etc for every steps in the down-sampling first, before the forward method?</p>

username: KFrank
cooked: <p>Hi Alan!</p>
<aside class="quote no-group" data-username="Alan_lam" data-post="3" data-topic="145623" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alan_lam/48/43715_2.png" class="avatar"> Alan_lam:</div>
<blockquote>
<p>I guess the <code>Transpose_Block</code> cannot be instantiatized inside forward method too because it otherwise would be re-created in every forward pass?</p>
</blockquote>
</aside>
<p>Yes, that is correct.</p>
<aside class="quote no-group">
<blockquote>
<p>what is the solution if I still want to pass the <code>output_size</code> argument when instantiatizing the <code>Transpose_Block</code>?</p>
</blockquote>
</aside>
<p>You do pass in an <code>output_size</code> argument when you instantiate<br>
<code>Transpose_Block</code> (but it doesn’t make sense to try to do it this way).</p>
<p>However, you do <em>not</em> pass in <code>output_size</code> when you instantiate the<br>
<code>ConvTranspose2d</code> <code>block</code> property of <code>Transpose_Block</code>. Instead,<br>
you pass in <code>output_size</code> as an additional argument when you apply<br>
the instance of <code>ConvTranspose2d</code> to <code>x</code> in <code>Transpose_Block</code>’s <code>forward()</code><br>
method.</p>
<p>So you don’t need to know <code>output_size</code> when you instantiate<br>
<code>Transpose_Block</code> and / or <code>ConvTranspose2d</code> – you just need to know<br>
it when you apply an already-instantiated <code>Transpose_Block</code> to <code>x</code> in the<br>
<code>for up</code> loop in <code>Generator</code>’s <code>forward()</code> method.</p>
<aside class="quote no-group">
<blockquote>
<p>Do I have to manually calculate the <code>size_list</code> according to the <code>kernel_size</code>, <code>padding</code> etc for every steps in the down-sampling first, before the forward method?</p>
</blockquote>
</aside>
<p>No, I don’t believe so.  I don’t fully understand your logic, but I believe<br>
that you can continue to compute the <code>self.list</code> of <code>output_size</code>s in<br>
the <code>for (index, down)</code> loop of <code>Generator</code>’s <code>forward()</code> method, and<br>
then pass those <code>output_size</code>s in as a newly-added argument to<br>
<code>Transpose_Block</code>’s <code>forward()</code> method which will then pass it along<br>
into the call where you apply that <code>Transpose_Block</code>’s<br>
<code>ConvTranspose2d</code> to <code>x</code>.</p>
<p>If you store <code>Generator</code>’s <code>Transpose_Block</code>s in a <code>ModuleList</code> that is<br>
a property of <code>Generator</code>, and each <code>ConvTranspose2d</code> is a property of<br>
its enclosing <code>Transpose_Block</code>, then when you call<br>
<code>model = Generator().to(DEVICE)</code> the <code>Transpose_Block</code>s and their<br>
contained <code>ConvTranspose2d</code>s should all be properly moved to the gpu<br>
(resolving your original error).</p>
<p>Best.</p>
<p>K. Frank</p>

username: Alan_lam
cooked: <p>Yes you are right. I got lost and confused mid-way in my thinking. I just need to separate the forward method of my <code>Transpose_Block</code> from others so that the <code>output_size</code> argument can be passed into it during the forward method of <code>Generator</code>. Thank you so much.</p>