What is the first argument in step() function?

[Jingnan-Jia]

I am new to ignite, I found several examples.

def train_step(trainer, batch):
    model.train()
    optimizer.zero_grad()
    x, y = prepare_batch(batch)
    y_pred = model(x)
    loss = loss_fn(y_pred, y)
    loss.backward()
    optimizer.step()
    return loss.item()

trainer = Engine(train_step)
trainer.run(data, max_epochs=100)

We can see the first argument in train_step() function is trainer. (I do not know why we need this parameter even we do not use it in the function ...)

def update(engine, batch):
    x, y = batch
    y_pred = model(inputs)
    loss = loss_fn(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item(), y_pred, y

trainer = Engine(update)

In this example, the first argument of update() function is engine . (I do not know why we need this parameter even we do not use it in the function ...)

I want to know:

1. why we need to add an argument trainer or engine in these function even we do not use them in the function?
2. whe trainer and engine are okay to be there even if they are not defined before?
3. Are there any other options apart from trainer and engine?

[ydcjeff]

Hi @Ordgod , you can provide the first argument with whatever name you want. The first argument is used here.

ignite/ignite/engine/engine.py

Line 852 in 5ceacbb

self.state.output = self._process_function(self, self.state.batch)

self._process_function is update or train_step as per your question which will do forward pass + calc loss + gradient update.

1. self itself is the Engine, so we use any class attribute, methods, etc of Engine whenever we need. So to answer your question, it's used internally.
2. Since trainer and engine are arguments, it's ok if they are not defined before.
3. The process_function only needs two arguments as per docs.

process_function (callable) – A function receiving a handle to the engine and the current batch in each iteration, and returns data to be stored in the engine’s state.

https://pytorch.org/ignite/master/engine.html#ignite.engine.engine.Engine

[Jingnan-Jia]

So the following code also works right? (use _ to replace trainer)

def update(_, batch):
    x, y = batch
    y_pred = model(inputs)
    loss = loss_fn(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item(), y_pred, y

trainer = Engine(update)

[ydcjeff]

Yes, it works

[sdesrozis]

Hi @Ordgod et very pleased to welcome you !!

As you can see, Engine is the main abstraction to loop on data and call a process function (ie train_step and update) on it. This function is defined by users. During the training, this process function is managed by the Engine and called in backend. You can see that the batch argument in the process function is the data but sometimes it is not sufficient. To add extensibility to the user, the first argument is the Engine itself. So you can catch everything about the current running flow (iterations, epochs, etc.).

HTH