Simple categorical cross entropy model not learning

kevs: <p>Hey, i have some training data in the format:<br>
X = [[0.1, 0.2], [0.1, 0.3], …] - (x,y) input data<br>
y = [[1,0,0], [0,1,0], … ] - one-hot encoded output</p>
<p>I have defined the model and training process as the following:</p>
<pre><code class="lang-auto">w0 = torch.tensor(np.random.randn(2,64), requires_grad=True)
w1 = torch.tensor(np.random.randn(64,3), requires_grad=True)

optim = torch.optim.Adam([w0, w1])

for epoch in range(1000):
        out = torch.tensor(X, requires_grad=True).matmul(w0).relu().matmul(w1)
        out = torch.nn.functional.log_softmax(out, dim=1)

        y_pred = out.detach().numpy().argmax(axis=1)
        acc = np.mean(y_pred == y_true) # y_true is the categorical representation

        # Categorical cross-entropy loss
        loss = -(y * out).sum(dim=1).mean()

        optim.zero_grad()
        loss.backward()
        optim.step()

        print(f"loss: {loss.data}, acc: {acc}")
</code></pre>
<p>The accuracy is stuck at 40% when it should be around 90%.</p>

KFrank: <p>Hi Kevin!</p>
<p>The short story is that I don’t see anything wrong with what you are doing.</p>
<p>Some further comments, below:</p>
<aside class="quote no-group quote-modified" data-username="kevs" data-post="1" data-topic="109506" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kevs/48/33660_2.png" class="avatar"> kevs:</div>
<blockquote>
<p>X = [[0.1, 0.2], [0.1, 0.3], …] - (x,y) input data<br>
y = [[1,0,0], [0,1,0], … ] - one-hot encoded output</p>
<pre><code class="lang-auto">w0 = torch.tensor(np.random.randn(2,64), requires_grad=True)
w1 = torch.tensor(np.random.randn(64,3), requires_grad=True)

optim = torch.optim.Adam([w0, w1])

for epoch in range(1000):
        out = torch.tensor(X, requires_grad=True).matmul(w0).relu().matmul(w1)
        out = torch.nn.functional.log_softmax(out, dim=1)
</code></pre>
</blockquote>
</aside>
<p>Note, you don’t need or want <code>X</code> to have <code>requires_grad=True</code>.  You’re<br>
not optimizing the values of <code>X</code> (and you haven’t added <code>X</code> to <code>optim</code>).<br>
However, I don’t think this actually hurts anything, so I don’t think it’s<br>
a problem.  Also, if <code>X</code> is already a pytorch tensor, you don’t need to<br>
wrap it in a newly-constructed tensor (not even for the purpose of turning<br>
on <code>requires_grad=True</code>).</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        y_pred = out.detach().numpy().argmax(axis=1)
        acc = np.mean(y_pred == y_true) # y_true is the categorical representation
</code></pre>
</blockquote>
</aside>
<p>This only involves calculating <code>acc</code>, which is not part of the <code>loss</code> you<br>
optimize, so this is not an issue.  However, I find it good practice to make<br>
a point of performing such calculations using pytorch tensor operations,<br>
rather than numpy, wherever possible.  If this were part of <code>loss</code>, the<br>
numpy calculations would “break the computation graph,” and your<br>
training would fail.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        # Categorical cross-entropy loss
        loss = -(y * out).sum(dim=1).mean()

        optim.zero_grad()
        loss.backward()
        optim.step()

        print(f"loss: {loss.data}, acc: {acc}")
</code></pre>
</blockquote>
</aside>
<p>This looks correct (although I haven’t tried your code).</p>
<aside class="quote no-group">
<blockquote>
<p>The accuracy is stuck at 40% when it should be around 90%.</p>
</blockquote>
</aside>
<p>Is there anything perniciously non-linear about the relationship between<br>
<code>X</code> and <code>y</code>?  It can be hard to fit some functions with a one-hidden-layer<br>
network.</p>
<p>Some suggestions to try:  Play around with the learning rate – lower,<br>
higher, or maybe use a learning-rate schedule (which you can easily<br>
do by hand, if you choose).</p>
<p>I also like to try plain-vanilla <code>SGD</code> (with and without momentum).  Even<br>
though things like <code>Adam</code> often train faster, they can be less robust.</p>
<p>It might be fun to re-implement your training using a <code>Model</code> with two<br>
<code>Linear</code> layers (without biases, to match what your current code is doing).<br>
Convert your one-hot targets to integer categorical labels using <code>argmax()</code><br>
(you can do this on-the-fly, if you want), and use <code>CrossEntropyLoss</code>.  This<br>
should be the same as what you are already doing, but if it works, then it<br>
means you have a bug somewhere that we haven’t caught, and gives you<br>
a cross-check that might help you track down the bug.</p>
<p>If you could tell us where your training data, <code>X</code> and <code>y</code>, come from or how<br>
they are  generated, we may have some thoughts on whether your training<br>
task is unusually difficult for some reason.</p>
<p>Good luck.</p>
<p>K. Frank</p>

kevs: <p>Thank you for the extensive answer. My problem was the missing bias unit. Rewriting the forward pass as the following made the model learn non-linear patterns:</p>
<pre><code class="lang-auto">...
b0 = torch.tensor(np.random.randn(64))
b1 = torch.tensor(np.random.randn(3))
...
out = torch.tensor(X, requires_grad=True).matmul(w0).add(b0).relu().matmul(w1).add(b1)
out = torch.nn.functional.log_softmax(out, dim=1)
</code></pre>