Bug in Data Parallel?

Rafael_R: <p>Hi,</p>
<p>I have a network  having two components (X, Y) which takes two inputs: A and B.</p>
<p>I wrap the network in data parallel. X returns full batch size output (how?) while Y returns 1/4th batch size output (this is expected).</p>
<p>Should I not nest subclasses of nn.Module while using Data parallel wrapper? Since, data parallel does support multiple inputs so this is unexpected (<a href="https://github.com/pytorch/pytorch/pull/794" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/794</a>).</p>
<p>I even tried separating the X, Y modules and wrapping them in data parallel separately, I still get the same error.</p>
<p>Any hints on what might be wrong?</p>

ptrblck: <p>Could you post a minimal code snippet to see, how you are using submodules inside your model and how you are applying <code>nn.DataParallel</code> to it?</p>

Rafael_R: <p>I am getting the same error in two ways:</p>
<pre><code class="lang-auto">&gt; class EncoderCNN(nn.Module):
&gt;     def __init__(self, embed_size, pre_trained_emb_size=2048):
&gt;         """Load the pretrained ResNet Features"""
&gt;         super(EncoderCNN, self).__init__()
&gt;         self.linear = nn.Linear(pre_trained_emb_size, embed_size)
&gt;         self.relu = nn.ReLU(embed_size)
&gt;         self.bn = nn.BatchNorm1d(embed_size, momentum=0.01)
&gt; 
&gt;     def forward(self, images):
&gt;         """Extracted feature vectors form input"""
&gt;     
&gt;         features = self.bn(self.relu(self.linear(images)))
&gt;         return features
&gt; 
&gt; class TransformerEncoder(nn.Module):
&gt;     def __init__(self, vocab_size, embed_size, d_model=300, nhead=6, dim_feedforward=512, dropout=0.1):
&gt;         super(TransformerEncoder, self).__init__()
&gt;         self.input_size = vocab_size
&gt;         self.embed_size = embed_size
&gt;         self.embedding = nn.Embedding(self.input_size, self.embed_size)
&gt;         # Encoder CNN
&gt;         self.encoder_cnn = EncoderCNN(embed_size=self.embed_size)
&gt;         self.transformer_encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward)
&gt; 
&gt;     def forward(self, input, input_lengths, images):
&gt;         input = input.reshape((input.shape[1], input.shape[0]))
&gt;         embedded = self.embedding(input)
&gt;         output = self.transformer_encoder_layer(embedded)
&gt;         image_encodings = self.encoder_cnn(images)
&gt;         return output, image_encodings
</code></pre>
<p>Even if I take the encoder cnn module out and wrap it in data parallel separately, I get the same error. the final returned tensors output, image_encodings have different batch sizes <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>

ptrblck: <p>I’m not sure, why you are reshaping the input in <code>TransformerEncoder</code>'s <code>forward</code> method.<br>
Could you explain this work flow as I think it might be related to this issue?</p>

Rafael_R: <p>Thanks, <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<ol>
<li>
<p>I am sending max_len x batch size x embedding_size input to transformer layer. This is the same api as in LSTM (without batch_first=True). Since, data parallel module expects the module’s input to have batch in first dimension, I am passing batch as first dimension and then reshaping it before passing it to transformer.</p>
</li>
<li>
<p>Same for images, batch_size x embedding_size input is passed to the module.</p>
</li>
</ol>
<p>Do you see anything wrong here?</p>

ptrblck: <p>Thanks for the information.<br>
In that case you might want to use <code>.permute</code>, as <code>.reshape</code> might interleave the data:</p>
<pre><code class="lang-python">x = torch.tensor([[0., 0.],
                  [1., 1.],
                  [2., 2.],
                  [3., 3.]])

print(x.reshape(x.size(1), x.size(0)))
&gt; tensor([[0., 0., 1., 1.],
          [2., 2., 3., 3.]])
print(x.permute(1, 0))
&gt; tensor([[0., 1., 2., 3.],
          [0., 1., 2., 3.]])
</code></pre>
<p>I don’t see any obvious errors. Could you add print statements inside the <code>forward</code> method, which will print the shape as well as the device of the input, output and each intermediate tensor?</p>

Rafael_R: <p>Thanks,</p>
<p>The print statements and the output on the console is:</p>
<pre><code class="lang-auto">    def forward(self, input, input_lengths, images):
        ptvsd.break_into_debugger()
        print("Input", input.shape, input.get_device())
        input = input.permute((1, 0))  #input.reshape((input.shape[1], input.shape[0]))
        print("Permuted Input", input.shape, input.get_device())
        embedded = self.embedding(input)
        print("Embedded", embedded.shape, embedded.get_device())
        #packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
        output = self.transformer_encoder_layer(embedded)
        print("Output", output.shape, output.get_device())
        #output, _ = torch.nn.utils.rnn.pad_packed_sequence(output)
        image_encodings = self.encoder_cnn(images)
        print("Image Encodings", image_encodings.shape, image_encodings.get_device())
        return output, image_encodings
</code></pre>
<p>And the output is:</p>
<pre><code class="lang-auto">Input torch.Size([64, 30]) 0
Input torch.Size([64, 30]) 1
Permuted Input torch.Size([30, 64]) 0
Input torch.Size([64, 30]) 2
Permuted Input torch.Size([30, 64]) 1
Permuted Input torch.Size([30, 64]) 2
Input torch.Size([64, 30]) 3
Embedded torch.Size([30, 64, 300]) 0
Embedded torch.Size([30, 64, 300]) 2
Embedded torch.Size([30, 64, 300]) 1
Output torch.Size([30, 64, 300]) 1
Output torch.Size([30, 64, 300]) 0
Output torch.Size([30, 64, 300]) 2
Permuted Input torch.Size([30, 64]) 3
Image Encodings torch.Size([64, 300]) 1
Image Encodings Image Encodingstorch.Size([64, 300]) 2
 torch.Size([64, 300]) 0
Embedded torch.Size([30, 64, 300]) 3
Output torch.Size([30, 64, 300]) 3
Image Encodings torch.Size([64, 300]) 3
</code></pre>
<p>And the returned tensors have the shape:</p>
<pre><code class="lang-auto">shape:torch.Size([120, 64, 300])
device:device(type='cuda', index=0)
</code></pre>
<p>and</p>
<pre><code class="lang-auto">shape:torch.Size([256, 300])
device:device(type='cuda', index=0)
</code></pre>

ptrblck: <p>It looks like you would need to permute the output of the <code>TransformerEncoderLayer</code> again, as its output has the shape <code>[T, N, E]</code>  (batch dimension in dim1).</p>

Rafael_R: <p>thanks, the data parallel issue is solved it seems and the code is working. however, it is very slow (slower than a single gpu ) and i am getting the error here: <a href="https://discuss.pytorch.org/t/how-to-flatten-parameters/53799" class="inline-onebox">How to flatten parameters?</a></p>
<p>any hints on how to fix the issue, i assume flatten lstm must be called after each call but then even if i do it, i still get the linked warning and the code is very slow</p>