'NoneType' object is not subscriptable when TensorboardLogger handler is EPOCH_COMPLETED

[Roy-Kid]

Hi, I have a weird issue. When I try to add an output handler with

    def attach_tensorboard(
        self,
        log_dir: str,
        event_name: EventEnum = Events.EPOCH_COMPLETED,
    ):
        tb_logger = TensorboardLogger(log_dir)
        tb_logger.attach_output_handler(
            self.trainer,
            event_name=event_name,
            tag="trainer",
            output_transform=lambda x: {"loss": x["loss"]},
            global_step_transform=global_step_from_engine(self.trainer),
        )

I got an error: 'NoneType' object is not subscriptable. It seems the state is reset before output. If I use ITERATION_COMPLETED, it works fine.

I tried both master and 0.5.1, and the engine is created with create_supervised_trainer.

Can someone help with this, give me a hint about where I made mistake.

[vfdev-5]

@Roy-Kid is it possible to have a small reproducer for the issue?
Also the full traceback can be helpful

[Roy-Kid]

Hi, this is how I did now, but not minimal reproducible snippet.

import molpot as mpot
import torch
from ignite.engine import Engine
from ignite.engine.events import EventEnum, Events, State
from ignite.metrics import MeanAbsoluteError, EpochWise, BatchWise
from pathlib import Path
from molpot import alias
from ignite.handlers import (
    ProgressBar,
    TensorboardLogger,
    global_step_from_engine,
)

# 1. get rMD17 dataset
rmd17_ds = mpot.dataset.rMD17(
    molecule="aspirin",
    save_dir="data",
    device="cpu",
    total=1000,
    processes=[mpot.process.NeighborList(cutoff=5.0)]
)

train_ds, eval_ds = torch.utils.data.random_split(rmd17_ds, [0.95, 0.05])
train_dl = mpot.DataLoader(train_ds, batch_size=1)
eval_dl = mpot.DataLoader(eval_ds, batch_size=1)

# 2. define model
pinet = mpot.potential.nnp.PiNet(
    depth=5,
    basis_fn=mpot.potential.nnp.radial.GaussianRBF(10, 5.0),
    cutoff_fn=mpot.potential.nnp.cutoff.CosineCutoff(5.0),
    pi_nodes=[64, 64],
    ii_nodes=[64, 64, 64, 64],
    pp_nodes=[64, 64, 64, 64],
    activation=torch.nn.Tanh(),
    rank=1,
)
e_readout = mpot.potential.nnp.readout.Atomwise(
    in_keys=[("pinet", "p1")],
    out_keys=[("predicts", "energy")],
    n_neurons=[64, 1],
    reduce="sum",
)
f_readout = mpot.potential.nnp.readout.Derivative(
    fx_key=("predicts", "energy"), dx_key=alias.pair_diff, out_keys=("predicts", "forces"), 
)
model = mpot.potential.PotentialSeq(pinet, e_readout)

# 3. training code
save_dir = Path("pinet2-rmd17")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)
loss_fn = mpot.engine.MultiTargetLoss(
    torch.nn.MSELoss(), [("energy", "energy", 1.0)]# , ("forces", "forces", 10.0)]
)

def _model_transform(output):
    return (output["predicts"], output["labels"])

def _train_output_transform(predicts, labels, loss):
    return {"predicts": predicts, "labels": labels, "loss": loss.item()}

def _eval_output_transform(predicts, labels):
    return {
        "predicts": predicts,
        "labels": labels,
    }

def update(
    engine, batch
):
    optimizer.zero_grad()
    model.train()
    outputs = model(batch)
    outputs = _model_transform(outputs)  # (predicts, labels)
    loss = loss_fn(*outputs)
    loss.backward()
    optimizer.step()
    return _train_output_transform(*outputs, loss)

trainer = Engine(update)

def get_metrics(pred_label, label):
    def _get_key(outputs):
        return (outputs["predicts"][pred_label], outputs["labels"][label])
    return _get_key

metric = MeanAbsoluteError(get_metrics("energy", "energy"))
metric.attach(trainer, "e_mae", BatchWise())

tb_logger = TensorboardLogger("tb_logs")
tb_logger.attach_output_handler(
    trainer,
    event_name=Events.EPOCH_COMPLETED,
    tag="trainer",
    output_transform=lambda x: {"loss": x["loss"]},
    global_step_transform=global_step_from_engine(trainer),
)
ProgressBar().attach(trainer)
trainer.run(
    train_dl, max_epochs=2
)

Error message:

Traceback (most recent call last):
  File "/workspaces/molcrafts/molpot/examples/test_tblogger.py", line 100, in <module>
    trainer.run(
  File "/opt/conda/lib/python3.12/site-packages/ignite/engine/engine.py", line 922, in run
    return self._internal_run()
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.12/site-packages/ignite/engine/engine.py", line 965, in _internal_run
    return next(self._internal_run_generator)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.12/site-packages/ignite/engine/engine.py", line 1040, in _internal_run_as_gen
    self._handle_exception(e)
  File "/opt/conda/lib/python3.12/site-packages/ignite/engine/engine.py", line 661, in _handle_exception
    raise e
  File "/opt/conda/lib/python3.12/site-packages/ignite/engine/engine.py", line 996, in _internal_run_as_gen
    self._fire_event(Events.EPOCH_COMPLETED)
  File "/opt/conda/lib/python3.12/site-packages/ignite/engine/engine.py", line 436, in _fire_event
    func(*first, *(event_args + others), **kwargs)
  File "/opt/conda/lib/python3.12/site-packages/ignite/handlers/tensorboard_logger.py", line 295, in __call__
    metrics = self._setup_output_metrics_state_attrs(engine, key_tuple=False)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.12/site-packages/ignite/handlers/base_logger.py", line 136, in _setup_output_metrics_state_attrs
    output_dict = self.output_transform(engine.state.output)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/molcrafts/molpot/examples/test_tblogger.py", line 96, in <lambda>
    output_transform=lambda x: {"loss": x["loss"]},
                                        ~^^^^^^^^
TypeError: 'NoneType' object is not subscriptable

If you need, I will replace my code with small MLP and random data.

[Roy-Kid]

Without my code, just use tensordict

import torch
from ignite.engine import Engine
from ignite.engine.events import EventEnum, Events, State
from ignite.metrics import MeanAbsoluteError, EpochWise, BatchWise
from pathlib import Path
from ignite.handlers import (
    ProgressBar,
    TensorboardLogger,
    global_step_from_engine,
)

from torchdata.nodes import IterableWrapper, ParallelMapper, Loader
from tensordict import TensorDict

node = IterableWrapper([
    TensorDict({
        "x": torch.randn(10),
        "labels": {"energy": torch.randn(10)},
        "predicts": {}
    }) for _ in range(10)
])
node = ParallelMapper(node, map_fn=lambda x: x**2, num_workers=3, method="thread")
train_dl = Loader(node)

class Model(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.kernel = torch.nn.Linear(10, 10)

    def forward(self, x):
        x["predicts", "energy"] = self.kernel(x["x"])
        return x

model = Model()

save_dir = Path("pinet2-rmd17")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)
mseloss_kernel = torch.nn.MSELoss()
loss_fn = lambda perdicts, labels: mseloss_kernel(perdicts["energy"], labels["energy"])

def _model_transform(output):
    return (output["predicts"], output["labels"])

def _train_output_transform(predicts, labels, loss):
    return {"predicts": predicts, "labels": labels, "loss": loss.item()}

def _eval_output_transform(predicts, labels):
    return {
        "predicts": predicts,
        "labels": labels,
    }

def update(
    engine, batch
):
    optimizer.zero_grad()
    model.train()
    outputs = model(batch)
    outputs = _model_transform(outputs)  # (predicts, labels)
    loss = loss_fn(*outputs)
    loss.backward()
    optimizer.step()
    return _train_output_transform(*outputs, loss)

trainer = Engine(update)

def get_metrics(pred_label, label):
    def _get_key(outputs):
        return (outputs["predicts"][pred_label], outputs["labels"][label])
    return _get_key

metric = MeanAbsoluteError(get_metrics("energy", "energy"))
metric.attach(trainer, "e_mae", BatchWise())

tb_logger = TensorboardLogger("tb_logs")
tb_logger.attach_output_handler(
    trainer,
    # event_name=Events.ITERATION_COMPLETED,  # OK
    event_name=Events.EPOCH_COMPLETED,  # not ok :(
    tag="trainer",
    output_transform=lambda x: {"loss": x["loss"]},
    global_step_transform=global_step_from_engine(trainer),
)
ProgressBar().attach(trainer)
trainer.run(
    train_dl, max_epochs=2
)

[vfdev-5]

Actually, I can not reproduce the issue, here is the code I'm running:

import torch
from ignite.engine import Engine
from ignite.engine.events import EventEnum, Events, State
from ignite.metrics import MeanAbsoluteError, EpochWise, BatchWise
from pathlib import Path
from ignite.handlers import (
    TensorboardLogger,
    global_step_from_engine,
)

model = torch.nn.Linear(10, 10)
train_dl = [torch.zeros(4, 10) for _ in range(10)]

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fn = torch.nn.MSELoss()


def _train_output_transform(predicts, labels, loss):
    return {"predicts": predicts, "labels": labels, "loss": loss.item()}

def _eval_output_transform(predicts, labels):
    return {
        "predicts": predicts,
        "labels": labels,
    }

def update(
    engine, batch
):
    optimizer.zero_grad()
    model.train()
    output = model(batch[0])
    outputs = (output, batch[1])
    loss = loss_fn(*outputs)
    loss.backward()
    optimizer.step()
    return _train_output_transform(*outputs, loss)

trainer = Engine(update)

tb_logger = TensorboardLogger("tb_logs")
tb_logger.attach_output_handler(
    trainer,
    event_name=Events.EPOCH_COMPLETED,  # MAYBE HERE YOU WOULD LIKE Events.ITERATION_COMPLETED TO LOG batch loss
    tag="trainer",
    output_transform=lambda x: {"loss": x["loss"]},
    global_step_transform=global_step_from_engine(trainer),
)
trainer.run(
    train_dl, max_epochs=2
)

[vfdev-5]

The problem is related to this line:

ignite/ignite/engine/engine.py

Line 1072 in 23f2a3f

self.state.batch = self.state.output = None

and the fact that the input train_dl does not have the length, so we accidentaly remove state.batch and state.output to estimate self.epoch_length.
To make it work you can run your code with specified epoch_length

trainer.run(
    train_dl, max_epochs=2, epoch_length=10,
)

This behaviour seems like a bug to me.

[Roy-Kid]

genius! It's indeed solved my problem! Thanks for much for your help!

[vfdev-5]

Glad to hear that it works for you! On the other hand if you want to plot the batch loss, maybe you would like to do that per iteration vs per epoch?

tb_logger.attach_output_handler(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    tag="trainer",
    output_transform=lambda x: {"loss": x["loss"]},
    global_step_transform=global_step_from_engine(trainer),
)

[Roy-Kid]

Thanks for your concern!

Our old TensorFlow code writes the loss or metrics to TensorBoard every N steps, which I believe is a better approach. Writing these values on every step could result in huge TensorBoard record files. I think using Events.ITERATION_COMPLETED(every=N) in Ignite would be an effective way to achieve this.

[vfdev-5]

Sounds good, Events.ITERATION_COMPLETED(every=N) should work as well. Feel free to open another question or issue ticket if something is not working as expected. We also have discord for instant messages.