Model doesn't train after custom weight initialization

username: artichoke
cooked: <p>I have a model that trains fine without weight initialization but when I try to initialize the weights of my model to custom values (generated by custom_weight_fn), it stops training. Here is the code:</p>
<pre><code class="lang-auto">def init_weights(m):
  if type(m) == nn.Linear or type(m) == nn.Conv2d:
    m.weight.data=custom_weight_fn(m.weight.shape)
    m.bias.data=custom_weight_fn(m.bias.shape)

model = nn.Sequential(
    nn.Conv2d(3, channel_1, 5, padding=2),
    nn.ReLU(),
    nn.Conv2d(channel_1, channel_2, 3, padding=1),
    nn.ReLU(),
    Flatten(),
    nn.Linear(32*32*channel_2, 10)
)

model.apply(init_weights)

</code></pre>
<p>Does anybody know why this is not working? The code runs without any errors but the model doesn’t train.</p>
<p>Thanks!</p>

username: ptrblck
cooked: <p>Don’t use the <code>.data</code> attribute as it’s deprecated and could yield any side effects.<br>
Instead, wrap the code in a <code>with torch.no_grad()</code> block and assign new <code>nn.Parameter</code>s to the module parameters or use <code>m.param.copy_</code> instead.</p>

username: artichoke
cooked: <p>Thanks for your suggestion. But when I apply these changes, it still doesn’t train… Here is my updated code:</p>
<pre><code class="lang-auto">def init_weights(m):
  with torch.no_grad():
    if type(m) == nn.Linear or type(m) == nn.Conv2d:
      m.weight = nn.Parameter(custom_weight_fn(m.weight.shape))
      m.bias = nn.Parameter(custom_weight_fn(m.bias.shape))

model = nn.Sequential(
    nn.Conv2d(3, channel_1, 5, padding=2),
    nn.ReLU(),
    nn.Conv2d(channel_1, channel_2, 3, padding=1),
    nn.ReLU(),
    Flatten(),
    nn.Linear(32*32*channel_2, 10)
)

model.apply(init_weights) #if I comment this line, the model trains with no problem
</code></pre>
<p>I also tried m.param.copy_ and unfortunately that doesn’t work either.</p>

username: ptrblck
cooked: <p>In that case I guess your custom init method might yield a bad initialization, as using e.g. <code>torch.randn</code> works fine and creates valid gradients:</p>
<pre><code class="lang-python">def custom_weight_fn(shape):
    return torch.randn(shape)


def init_weights(m):
  with torch.no_grad():
    if type(m) == nn.Linear or type(m) == nn.Conv2d:
      m.weight = nn.Parameter(custom_weight_fn(m.weight.shape))
      m.bias = nn.Parameter(custom_weight_fn(m.bias.shape))

channel_1 = 3
channel_2 = 3
model = nn.Sequential(
    nn.Conv2d(3, channel_1, 5, padding=2),
    nn.ReLU(),
    nn.Conv2d(channel_1, channel_2, 3, padding=1),
    nn.ReLU(),
    nn.Flatten(),
    nn.Linear(32*32*channel_2, 10)
)

model.apply(init_weights) 

out = model(torch.randn(1, 3, 32, 32))
out.mean().backward()

for name, param in model.named_parameters():
    print('{}, grad.abs().max(): {}'.format(name, param.grad.abs().max()))
</code></pre>

username: artichoke
cooked: <p>You are right! I had a small issue, somehow when I initialize my weights according to the custom init function, the training becomes much harder, but if I do hyper-parameter tuning then the model starts to train. Your suggestion on how to set the weights to custom values works well.</p>