Functional Conv2 accept a batch of weights

username: vpolflie
cooked: <p>Hi,</p>
<p>At the moment I am trying to implement a meta-learning algorithms and the size of the model is quite large so I am also trying  to use DataParallel. However, I am currently encountering an issue with one GPU taking the brunt of the load and running out of memory. This is since I generate weights for each sample in my batch, which means I have to loop over these weights and apply a functional conv and thus this operation cant be data paralleled and it ends up on the same GPU.</p>
<p>Is there any easy way to feed a batch of weights to a functional conv or are there any plans to implement this in pytorch in the near future?</p>
<p>Cheers,<br>
Vincent</p>

username: mrshenli
cooked: <aside class="quote no-group" data-username="vpolflie" data-post="1" data-topic="74530">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/v/f1d935/48.png" class="avatar"> vpolflie:</div>
<blockquote>
<p>Is there any easy way to feed a batch of weights to a functional conv</p>
</blockquote>
</aside>
<p>Not sure if I understand the request clearly, it would be helpful if you could share some pseudo code.</p>
<p>If all you need is scatter conv weights (one weight per sample) across different GPUs, looks like you can wrap that (samples + one conv layer per GPU) into a custom function? In the forward function, you can do sth like:</p>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, samples):
    outputs = []
    for sample in samples:
        weight = generate_per_sample_weight(sample)
        replace_conv_weight(self.conv, weight)
        outputs.append(self.conv(sample))
    return outputs
</code></pre>

username: vpolflie
cooked: <p>Sorry for the late response.</p>
<p>The code you provide is the pseudo code I would give and very similar to the code I have in my code base, one small change is that the weights are generated by different samples.</p>
<pre><code class="lang-auto">def forward(self, samples, reference_samples):
    outputs = []
    for sample in samples:
        weight = generate_per_sample_weight(reference_samples)
        replace_conv_weight(self.conv, weight)
        outputs.append(self.conv(sample))
    return outputs
</code></pre>
<p>However, the main issue with this is that pytorch doesn’t distribute this over multiple GPU’s because of the for loop. These calculation are all located on the first main GPU (with the standard dataparallel package)</p>
<p>I was wondering if there is an easy way to write something like this which still allows for data parallelisation:</p>
<pre><code class="lang-auto">weights: Batch x # INPUT FILTERS x # OUTPUT FILTERS x FILTER WIDTH x FILTER HEIGHT
samples: BATCH x CHANNELS x WIDTH x HEIGHT
def forward(self, samples, weights):
    outputs = self.conv(sample, weights)
    return outputs
</code></pre>
<p>So this self.conv function would then be one purely based on matrices like the original conv one, which should allow data parallelisation.</p>

username: mrshenli
cooked: <aside class="quote no-group" data-username="vpolflie" data-post="3" data-topic="74530">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/v/f1d935/48.png" class="avatar"> vpolflie:</div>
<blockquote>
<p>I was wondering if there is an easy way to write something like this which still allows for data parallelisation:</p>
</blockquote>
</aside>
<p>This should work, as <a href="https://github.com/pytorch/pytorch/blob/df8d6eeb19423848b20cd727bc4a728337b73829/torch/nn/parallel/data_parallel.py#L151-L155" rel="noopener nofollow ugc">DataParallel</a> simply replicates model and scatters input. (assuming <code>self.conv</code> is a customized conv layer that replaces weight) So if you wrap that with DataParallel, different thread/replica should see samples/weights on a different device. Did you encounter any issue when doing this?</p>

username: vpolflie
cooked: <p>At the moment I have the following psuedo code:</p>
<pre><code class="lang-auto">def batch_conv(x, weight, bias=None, stride=1):
    for i in range(x.size()[0]):
        yi = F.conv_transpose2d(x[i:i+1], weight=weight[i], bias=bias[i,:weight.size(2)], padding=1, stride=int(1/stride), output_padding=1, groups=groups)
        y = concat(y, yi)
    return y

class AdaptiveConv2d(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()

    def forward(self, input, weight=None, bias=None, stride=1):
        return batch_conv(input, weight, bias, stride)

</code></pre>
<p>However this doesn’t distribute properly and my assumption is that the data parallel isn’t able to handle the for loop in my code.</p>
<p>I haven’t tried implementing a conv layer that takes a batch of weights instead of a single sample. Since I have only switched to pytorch recently and I am a bit out of depth with this.</p>