Training a model with 2 models is not working

jin8495: <p>I am trying to combine 2 models sequentially like one’s output is another’s input.<br>
But second one which creates final output is pre-trained model and it will be trained no more. I only wants to train first one.</p>
<p>It is easier to understand with the picture.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/8/68c8670087560a897879a78ec2842de26c475ece.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/68c8670087560a897879a78ec2842de26c475ece" title="캡처"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/8/68c8670087560a897879a78ec2842de26c475ece_2_494x374.jpeg" alt="캡처" data-base62-sha1="eWX13i4dUh7rvZDQMG49GnYnmSW" width="494" height="374" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/8/68c8670087560a897879a78ec2842de26c475ece_2_494x374.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/8/68c8670087560a897879a78ec2842de26c475ece_2_741x561.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/8/68c8670087560a897879a78ec2842de26c475ece_2_988x748.jpeg 2x" data-dominant-color="D7C4C0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">캡처</span><span class="informations">1221×926 42.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>One model has architecture below,</p>
<pre><code class="lang-auto">class LinearModel(nn.Module):
    def __init__(self, bias=True):
        super(LinearModel, self).__init__()
        self.bias = bias
        
        self.linear0 = nn.Linear(64, 64, bias=bias)
        self.quantize = QuantizeLayer(num_bits=8)
        
    def forward(self, x):
        linear0_out = self.linear0(x)
        quant_out = self.quantize(linear0_out)
        return quant_out
</code></pre>
<p>And it has a ‘QuantizeLayer’ which I have made.</p>
<pre><code class="lang-auto">class Quantize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, scale, quant_min, quant_max):
        rounded_input = torch.round(input*scale) / scale
        return torch.clamp(rounded_input, min=quant_min, max=quant_max)
    
    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()
        return grad_input, None, None, None
        
class QuantizeLayer(nn.Module):
    def __init__(self, num_bits=8):
        super(QuantizeLayer, self).__init__()
        self.scale = 2 ** num_bits
        self.quant_min = -1
        self.quant_max = 1 - (2**-(num_bits-1))
        
        self.quantize = Quantize.apply
        
    def forward(self, x):
        return self.quantize(x, self.scale, self.quant_min, self.quant_max)
</code></pre>
<p>The other model has architecture below,</p>
<pre><code class="lang-auto">class LSTMModel(nn.Module):
    def __init__(self):
        super(LSTMModel, self).__init__()
        
        input_size=1
        hidden_size=256
        num_layers=1
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, 1)
        
    def forward(self, x):
        reshaped = torch.reshape(x, (-1, 64, 1))
        lstm_out, _ = self.lstm(reshaped)
        lstm_out = lstm_o[:,-1]
        linear_out = self.linear(lstm_out)
        
        return linear_out
</code></pre>
<p>And I wrapped up those two models like below,</p>
<pre><code class="lang-auto">class CombinedModel(nn.Module):
    def __init__(self, model0, model1):
        super(CombinedModel, self).__init__()
        self.model0 = model0
        
        for param in model1.parameters():
            param.requires_grad = False
        self.model1 = model1
        
    def forward(self, x):
        x = self.model0(x)
        x = self.model1(x)
        return x
</code></pre>
<p>Finally, I trained the model with dummy data to find out it is working or not.<br>
To check out, I compared all parameters before training and after training.</p>
<pre><code class="lang-auto"># Dummy data for data, labels
x = torch.rand((1000, 64))
y = torch.rand((1000, 1))

# Model objects
linearModel = LinearModel()
lstmModel = LSTMModel()
combinedModel = CombinedModel(linearModel, lstmModel)

# Parameter before training
state_dict_before = {}
for layer in combinedModel.state_dict():
    state_dict_before[layer] = combinedModel.state_dict()[layer].numpy()

# Training
criterion = nn.MSELoss()
optimizer = optim.SGD(combinedModel.parameters(), lr=1e-2)

combinedModel.train()
for i in range(10):
    pred = combinedModel(x)
    loss = criterion(pred, y)

    # backprop
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Parameter after training
state_dict_after = {}
for layer in combinedModel.state_dict():
    state_dict_after[layer] = combinedModel.state_dict()[layer].numpy()

# Compare parameter before training and after training
for layer in state_dict_before:
    print(np.equal(state_dict_before[layer], state_dict_after[layer]).all())
</code></pre>
<p>And all parameters are the same.</p>
<p>So I checked that gradients are really passing through the model while training by printing out “grad_out” and “grad_in” from ‘Quantize::backward()’. And it was not all-zero tensor.</p>
<p>I am digging whole day try to find out what went wrong. But I couldn’t figured out. ;(</p>

jin8495: <p>Ok. I have figured it out.</p>
<p>‘state_dict_before’ and ‘state_dict_after’ has pointers of parameters. It was not deep-copied when I assigned them.</p>