Simple custom optimizer in a NN via model.parameters()?

username: Michi01
cooked: <p>Given a “standard” NN in PyTorch:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy as np
import os
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from AugustVonDezent import Adaam

learning_rate = 0.01
BATCH_SIZE = 64



device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()
loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.03)

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()

        loss.backward()
        
        for param in model.parameters():
            param.

        optimizer.step()








        #optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss: &gt;7f} [{current: &gt;5d}/{size: &gt;5d}]")

def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X,y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):&gt;0.1f}%, Avg loss: {test_loss:&gt;8f} \n")




epochs = 2
# Ultra ELITE optimizer
for t in range(epochs):
    optimizer = Adaam(model.parameters(), lr=0.35)

    print(f"Epoch {t + 1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
    print("Done!")

#Standard optimizer
for t in range(epochs):
    optimizer = torch.optim.SGD(model.parameters(), lr=0.35)
    print(f"Epoch {t + 1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
</code></pre>
<p>i know that one can access the param by doing a for loop in the model itself:</p>
<p>for param in model.parameters():<br>
param.grade = None (same as optimizer.zero_grad)</p>
<p>However, is it possible to adjust the weights in such a for loop so that one can implement a “simple” custom optimizer without touching the optim library?</p>
<p>For example, im looking for something like this:<br>
for param in model.parameters():<br>
param = param - learning_rate * loss</p>
<p>I know i can create a custom optimizer following the optim optimizer examples (already did that), but i would like to add the backpropogation and adjust the weights as shown above. Does anyone have here a quick tip on how to do that? Thanks in advance</p>

username: JuanFMontesinos
cooked: <p>I mean, no one forces you to use optimizers.<br>
In the end optimizers do a for loop as well, the same way but adding more complex mechanisms. You can do  a loop exactly and use model.zero_grad to reset the gradients or set the gradients directly to None.</p>

username: Michi01
cooked: <p>Hey! Thanks for your answer and taking time!</p>
<p>Right, but how do i update the weights? If i remove in my code optimizer.zero_grad() and also optimizer.step() and add instead the for loop like so:</p>
<pre><code class="lang-auto">     for param in model.parameters():
            param.grad = None
            param.grad = -learning_rate * loss
</code></pre>
<p>i get RuntimeError: assigned grad has data of a different size as an error. I am doing something wrong? I’ve read this is related due to the device im using (cpu vs gpu) but as far as i can tell i don’t declare anywhere which device im using for the NN. Thanks again for your time!</p>

username: AlphaBetaGamma96
cooked: <aside class="quote no-group" data-username="Michi01" data-post="1" data-topic="151204">
<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/michi01/48/47149_2.png" class="avatar"> Michi01:</div>
<blockquote>
<p>param = param - learning_rate * loss</p>
</blockquote>
</aside>
<p>This should be,</p>
<pre><code class="lang-auto">param = param - learning_rate * grad(loss)
</code></pre>
<p>It also explains why you have the mismatch error as <code>learning_rate</code> and <code>loss</code> are both scalars and your weights will be a matrix.</p>
<aside class="quote no-group" data-username="Michi01" data-post="3" data-topic="151204">
<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/michi01/48/47149_2.png" class="avatar"> Michi01:</div>
<blockquote>
<p>i get RuntimeError: assigned grad has data of a different size as an error.</p>
</blockquote>
</aside>
<p>Also, could you explain why you don’t want to use the <code>torch.optim.Optimizer</code> class? It exists to update your weights efficiently and scalably.</p>

username: Michi01
cooked: <p>Its for a presentation, the implementation should be simple and easy to understand. I feel using the optim implementation makes it a bit bulky for what im trying to do.</p>
<p>I have tried</p>
<pre><code class="lang-auto">param = param - learning_rate * loss
</code></pre>
<p>before, but i doesn’t really affect the loss. What is the grad function you are calling or is just a typo? Because your code says:</p>
<pre><code class="lang-auto">param = param - learning_rate * grad(loss)
</code></pre>
<p>Thanks for taking your time! Appreciate it!</p>

username: JuanFMontesinos
cooked: <p>you need to use param.grad which is the the derivative of the loss/parameter so it would be<br>
<code>param = param - learning_rate * param.grad</code></p>

username: AlphaBetaGamma96
cooked: <p>As <a class="mention" href="/u/juanfmontesinos">@JuanFMontesinos</a> stated you need to use,<br>
<code>param = param - learning_rate * param.grad</code></p>
<p>The <code>grad</code> I had was just pseudo-code to say you need the gradient of the loss and not the loss value itself when performing gradient descent. When you call <code>loss.backward()</code> autograd computes the gradient via reverse-mode AD and stores the gradient as the <code>.grad</code> attribute of the <code>param</code>.</p>

username: Michi01
cooked: <p>Thanks for your help both of you!</p>
<p>That clarified a lot, one last thing: I have tried your suggestion before and i always get this error:</p>
<pre><code class="lang-auto">   param = param - learning_rate * param.grad
TypeError: unsupported operand type(s) for *: 'float' and 'NoneType'
</code></pre>
<p>Looking up the internet, people suggests to use torch.no_grad(): before updating the weight, so i’ve tried different things out, for example:</p>
<pre><code class="lang-auto">
        for param in model.parameters():
            with torch.no_grad():
                param -= learning_rate * param.grad
</code></pre>
<p>however, the error still exists. Any ideas why that is so?</p>

username: AlphaBetaGamma96
cooked: <aside class="quote no-group" data-username="Michi01" data-post="3" data-topic="151204">
<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/michi01/48/47149_2.png" class="avatar"> Michi01:</div>
<blockquote>
<pre><code class="lang-auto">            param.grad = None
</code></pre>
</blockquote>
</aside>
<p>Do you still have this line? If so, remove it. (And check you’ve called <code>loss.backward()</code> beforehand).</p>
<p>Your current <code>param.grad</code> attribute is <code>None</code>. So, you’ve either overwritten it with <code>None</code> or you haven’t called <code>loss.backward()</code></p>

username: JuanFMontesinos
cooked: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="9" data-topic="151204">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p><code>None</code></p>
</blockquote>
</aside>
<p>If you find None it either means you didn’t call loss.backward() or that your backpropagation is broken do to a bug.<br>
So the pseudo code would be</p>
<pre><code class="lang-auto">
for param in parameters:
  param.grad= None
out = forward(inputs)
loss = criterion(out,gt)
loss.backward()
for param in parameters:
    param = param - lr * param.grad
</code></pre>

username: Michi01
cooked: <p>Thanks for the help everyone!</p>
<p>I was able to find now a working solution and post it here, just in case somebody else runs into that same problem. Given the code base from above, i changed the following:</p>
<pre><code class="lang-auto">
def update_function(param, grad, loss, learning_rate):
  return param - learning_rate * grad

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):


        pred = model(X)

        model.zero_grad()
        loss = loss_fn(pred, y)

        loss.backward()
        with torch.no_grad():
            for x in model.parameters():
                updated_val = update_function(x, x.grad, loss, learning_rate)
                x.copy_(updated_val)



        #optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss: &gt;7f} [{current: &gt;5d}/{size: &gt;5d}]")
</code></pre>
<p>Works now fine and the weights are now being updated correctly!</p>
<p>Thanks and have a great one everyone!</p>