Different prediction results, update different weights in the model

shirui-japina: <p>Suppose there are <em><strong>parts A</strong></em> and <em><strong>part B</strong></em> in a model.</p>
<p>What I want to achieve is:</p>
<ul>
<li>
<p>If the model prediction is <strong>a</strong>, then update the weight of <em><strong>part</strong> A</em> in the model.</p>
</li>
<li>
<p>If the model prediction is <strong>b</strong>, then update the weight of <em><strong>part</strong> B</em> in the model.</p>
</li>
</ul>
<p>My question is:</p>
<ol>
<li>How is it implemented in PyTorch? (Especially when batch prediction)</li>
<li>From the perspective of deep learning, is such an update scheme practical?</li>
</ol>

vainaijr: <p>do you mean something like this,</p>
<pre><code class="lang-auto">import torch.nn as nn, torch

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.A = nn.ModuleDict([['lin', nn.Linear(3, 3)], 
                             ['dropout', nn.Dropout(0.5)]])
    self.B = nn.Linear(3, 3)
  def forward(self, x):
    out = self.B(x)
    out = self.A['lin'](out)
    out = self.A['dropout'](out)
    return out

net = Net()
</code></pre>
<pre><code class="lang-auto">optimizer1 = torch.optim.SGD(net.A.parameters(), lr=0.01)
optimizer2 = torch.optim.SGD(net.B.parameters(), lr=0.01)
</code></pre>
<pre><code class="lang-auto">for i in range(2):
  optimizer1.zero_grad()
  optimizer2.zero_grad()
  input = torch.randn(3, 3)
  loss = (net(input)).sum()
  loss.backward()
  if loss &lt; 1: # replace with net(input) == prediction
    optimizer1.step()
  else:
    optimizer2.step()
</code></pre>
<p>or with one optimizer</p>
<pre><code class="lang-auto">optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
</code></pre>
<pre><code class="lang-auto">for i in range(2):
  optimizer.zero_grad()
  input = torch.randn(3, 3)
  loss = (net(input)).sum()
  print(loss)
  loss.backward()
  if loss &lt; 1: # replace with net(input) == prediction
    net.A.lin.requires_grad_(False)
    net.B.requires_grad_(True)
  else:
    net.A.lin.requires_grad_(True)
    net.B.requires_grad_(False)
  print(list(net.parameters()))
  optimizer.step()
</code></pre>

shirui-japina: <p>Thanks for your reply.<img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:"><img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:"><br>
The code you wrote is very good.<img src="https://discuss.pytorch.org/images/emoji/apple/tada.png?v=9" title=":tada:" class="emoji" alt=":tada:"><img src="https://discuss.pytorch.org/images/emoji/apple/tada.png?v=9" title=":tada:" class="emoji" alt=":tada:"><img src="https://discuss.pytorch.org/images/emoji/apple/tada.png?v=9" title=":tada:" class="emoji" alt=":tada:"></p>
<p>But I have a question: When the optimizer is BGD, MBGD, can the above code run normally?<img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=9" title=":thinking:" class="emoji" alt=":thinking:"><img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=9" title=":thinking:" class="emoji" alt=":thinking:"></p>
<p>In the batches, there will be samples predicted as <strong>a</strong>, and samples predicted as <strong>b</strong>.<br>
I don’t think <code>optimizer.step()</code> can run normally at that time.<img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=9" title=":thinking:" class="emoji" alt=":thinking:"><img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=9" title=":thinking:" class="emoji" alt=":thinking:"></p>

vainaijr: <p>but if we make predictions for a batch, then optimizer.step() would only update parameters once, for example, if our predictions were something like [a, b, a, a, b], and then we do optimizer.step(), then it would update parameters of our neural network once, if we want to update after every prediction model makes, then we will have to do optimizer.step() after every prediction, so model predicts something like [a], we do optimizer.step() on some parameters, model predicts [b], we do optimizer.step() on some different parameters.</p>

shirui-japina: <p>So… that ‘Different prediction results, update different weights in the model’ can be implemented only by setting <strong>SGD</strong> as the optimizer? Since predictions like [a, b, a, a, b] would update different part of model (<em><strong>part A</strong></em> or <em><strong>part B</strong></em>) due to the different elements of that.</p>
<p>Is my understanding right?<img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=9" title=":thinking:" class="emoji" alt=":thinking:"><img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=9" title=":thinking:" class="emoji" alt=":thinking:"></p>

vainaijr: <p>let us consider one example, suppose we have a last word prediction task,</p>
<pre><code class="lang-auto"># I ate an apple
# 0  1  2   3

# I went to park
# 0  4   5   6

# I slept all day
# 0   7    8   9
</code></pre>
<p>so, our input to model will be like,</p>
<pre><code class="lang-auto">input = torch.LongTensor([[0, 1, 2], 
                      [0, 4, 5], 
                      [0, 7, 8]])
</code></pre>
<p>and target would be like,</p>
<pre><code class="lang-auto">target = torch.LongTensor([[3], 
                       [6], 
                       [9]])
# means if model sees 'I', 'ate', 'an' then it should predict 'apple'
# for 'I', 'went', 'to' predict 'park' and so on
</code></pre>
<p>now we create our dataset,</p>
<pre><code class="lang-auto">tensor_dataset = torch.utils.data.TensorDataset(input, target)
list(tensor_dataset)
</code></pre>
<p>[(tensor([0, 1, 2]), tensor([3])),<br>
(tensor([0, 4, 5]), tensor([6])),<br>
(tensor([0, 7, 8]), tensor([9]))]</p>
<p>let us consider a batch size of 2, so in the first batch, first two sentences would be considered, in the second batch, third sentence would be considered.</p>
<pre><code class="lang-auto">dataset = torch.utils.data.DataLoader(tensor_dataset, batch_size=2)
</code></pre>
<p>this is what our dataset looks like</p>
<pre><code class="lang-auto">for i, (input, target) in enumerate(dataset):
  print('i', i, '\ninput', input, '\ntarget', target)
</code></pre>
<pre><code class="lang-auto">i 0 
input tensor([[0, 1, 2],
        [0, 4, 5]]) 
target tensor([[3],
        [6]])
i 1 
input tensor([[0, 7, 8]]) 
target tensor([[9]])
</code></pre>
<p>now, we create our model,</p>
<pre><code class="lang-auto">class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.embed = nn.Embedding(10, 20, sparse=True) # our vocabulary has 10 words, we set embedding size as 20
    # we use sparse=True, so that embeddings of words that are not in a batch would not be updated
    self.transformer_encoder_layer = nn.TransformerEncoderLayer(20, 2)
    self.transformer_encoder = nn.TransformerEncoder(self.transformer_encoder_layer, 1)
    self.lin = nn.Linear(3*20, 10) # we get a probability distribution for all 10 words
    self.softmax = nn.Softmax(dim=-1) # just to print probabilities
  def forward(self, input):
    embedded_input = self.embed(input)
    print(input.size(0), input.size(1))
    print('embedded_input.shape', embedded_input.shape)
    out = self.transformer_encoder(embedded_input)
    out = out.reshape(out.size(0), out.size(1)*out.size(2))
    out = self.lin(out)
    print('out.shape', out.shape)
    print(self.softmax(out))
    return out
</code></pre>
<p>now, we do our training,</p>
<pre><code class="lang-auto">net = Net()
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
loss_plot = []

for num_epochs in range(10):
  for i, (input, target) in enumerate(dataset):
    optimizer.zero_grad()
    loss = loss_fn(net(input), target.view(-1))
    loss_plot.append(loss)
    loss.backward()
    for name, param in net.named_parameters():
      if 'embed' in name:
        print(name, param.grad)
    # print(list(net.parameters()))
    optimizer.step()
</code></pre>
<p>output of print statements looks something like this,</p>
<pre><code class="lang-auto">2 3
embedded_input.shape torch.Size([2, 3, 20])
out.shape torch.Size([2, 10])
tensor([[0.0136, 0.0095, 0.0165, 0.7450, 0.0182, 0.0222, 0.0389, 0.0225, 0.0256,
         0.0880],
        [0.0246, 0.0211, 0.0162, 0.0714, 0.0065, 0.0115, 0.7557, 0.0202, 0.0166,
         0.0562]], grad_fn=&lt;SoftmaxBackward&gt;)
embed.weight tensor(indices=tensor([[0, 1, 2, 0, 4, 5]]),
       values=tensor([[-1.2932e-02,  1.0104e-03, -3.9477e-03,  3.8568e-03,
                       -9.8572e-03,  1.3827e-03, -1.4141e-02,  7.4199e-03,
                        9.7238e-04, -6.7569e-03,  6.8323e-03,  8.7705e-03,
                       -6.8232e-04,  6.6810e-03,  8.1223e-03, -1.0538e-03,
                       -1.1698e-02,  8.6740e-03,  1.7381e-03, -5.2466e-03],
                      [ 6.2502e-05, -3.4446e-03, -8.7784e-03, -2.4090e-03,
                        3.3907e-03,  1.0026e-04,  4.1565e-03, -4.5490e-03,
                        5.8510e-03, -5.8728e-03, -7.7380e-03, -4.1700e-03,
                       -3.1775e-03,  6.2709e-03,  4.3878e-03,  1.1901e-03,
                       -9.7035e-04,  4.6694e-04,  1.4419e-03,  7.2143e-03],
                      [ 2.1315e-03, -2.1100e-03, -4.7602e-03, -3.3467e-03,
                       -3.1354e-03,  6.1653e-03, -1.1522e-02,  2.0889e-03,
                       -4.1310e-03,  1.0974e-03, -3.4026e-03, -3.2140e-03,
                       -6.5089e-03,  3.2289e-03,  8.5952e-04,  3.0189e-03,
                        5.5930e-03,  7.7494e-04,  4.9238e-03,  7.5434e-03],
                      [ 1.6089e-03, -1.5056e-02,  1.1222e-02,  3.6189e-03,
                        4.0124e-03, -6.9139e-03, -1.7996e-03,  2.3435e-03,
                        2.9756e-03,  4.3685e-03, -1.4790e-02,  2.6396e-03,
                       -4.5901e-03, -4.5700e-03, -5.2482e-03,  3.9613e-03,
                        4.8062e-03, -4.5896e-03,  1.2451e-02, -7.3047e-03],
                      [-7.6837e-03,  7.0184e-03,  5.9852e-03, -2.7696e-03,
                       -1.2842e-02, -2.8121e-03, -3.3888e-03,  7.1801e-04,
                       -1.1498e-02,  1.0984e-03, -1.7498e-03,  3.7961e-03,
                        1.8652e-03,  7.1568e-03, -4.6547e-03,  3.0268e-03,
                        9.9045e-03, -3.4706e-03, -3.0367e-03,  5.5819e-03],
                      [-1.8114e-02, -6.5709e-03,  5.3046e-03,  5.8249e-03,
                        9.4378e-04,  1.2052e-02,  9.1222e-03,  4.8452e-04,
                       -7.5793e-03, -3.2206e-03,  6.3282e-03, -1.7706e-02,
                        1.6900e-02, -2.1886e-03,  6.7194e-04, -4.2445e-03,
                       -3.8267e-04,  7.8978e-03, -9.4713e-03, -7.2790e-03]]),
       size=(10, 20), nnz=6, layout=torch.sparse_coo)
1 3
embedded_input.shape torch.Size([1, 3, 20])
out.shape torch.Size([1, 10])
tensor([[0.0180, 0.0087, 0.0069, 0.0273, 0.0210, 0.0133, 0.0216, 0.0091, 0.0049,
         0.8693]], grad_fn=&lt;SoftmaxBackward&gt;)
embed.weight tensor(indices=tensor([[0, 7, 8]]),
       values=tensor([[ 0.0210,  0.0100, -0.0089, -0.0045,  0.0082, -0.0014,
                        0.0132, -0.0120, -0.0027, -0.0058,  0.0039, -0.0031,
                       -0.0098, -0.0070, -0.0011,  0.0032, -0.0043,  0.0151,
                        0.0008, -0.0038],
                      [ 0.0050,  0.0010, -0.0012, -0.0027, -0.0032,  0.0006,
                        0.0040, -0.0026, -0.0033,  0.0074, -0.0001, -0.0050,
                        0.0011,  0.0009,  0.0014,  0.0017, -0.0023, -0.0074,
                        0.0033,  0.0004],
                      [ 0.0003,  0.0092,  0.0041,  0.0042,  0.0023, -0.0032,
                        0.0008, -0.0046, -0.0023,  0.0003, -0.0032,  0.0084,
                       -0.0021, -0.0017, -0.0119,  0.0003, -0.0076, -0.0019,
                        0.0012,  0.0016]]),
       size=(10, 20), nnz=3, layout=torch.sparse_coo)
</code></pre>
<p>here, we update embeddings of only inputs that are passed to our model, and not all the inputs, we have batch size of 2, so when we pass,</p>
<pre><code class="lang-auto">I ate an 
I went to
</code></pre>
<p>so embeddings of these words gets updated and embeddings of ‘slept’, ‘all’, ‘day’, ‘park’, ‘apple’ do not get updated.</p>
<p>when we pass</p>
<pre><code class="lang-auto">I slept all
</code></pre>
<p>as input, then embeddings of ‘ate’, ‘an’, ‘went’, ‘to’, ‘day’, ‘park’, ‘apple’ do not get updated.</p>
<p>if we want to update different embeddings based on what model predicted, that is, if model predicts</p>
<pre><code class="lang-auto">I ate an park
</code></pre>
<p>then we would want to update embeddings of some words, and if model predicted</p>
<pre><code class="lang-auto">I ate an apple
</code></pre>
<p>then we would want to update embeddings of some different words, then we would have to change this sparse tensor’s requires_grad of some indices, after model prediction, which currently I do not know how to do.</p>

shirui-japina: <p>Thank you so much… <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=9" title=":sob:" class="emoji" alt=":sob:"> <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=9" title=":sob:" class="emoji" alt=":sob:"></p>