Side effects when loading in lodeSTAR checkpoint

[naiikaa]

Hello! I just tried out the multi cell tracking with the deeplay lodeSTAR implementation following the notebook. Everything worked fine and I got the following satisfying result when pushing an image through after training:

Afterwards I wanted to try out saving and loading a checkpoint and performing the same plot. Since you use a lightning trainer these lines should be fine to load in the checkpoint after training:

loaded = dl.LodeSTAR.load_from_checkpoint(<.cpkt-path>, n_transforms=4,  optimizer=dl.Adam(lr=1e-4))

loaded.model.eval()

When I now run the same plot after loading in the model I get this result. I think it's still right but some form of translation is happening. I think loading in the model this way might be wrong on my side.

Thanks a lot for your help!

[BenjaminMidtvedt]

@naikaida Hi! I think you basically got it right. Did you try changing

lodestar.model.eval() 

to

lodestar.eval()

[naiikaa]

@BenjaminMidtvedt Thanks for your reply. I tried that also, but it resulted in the same behavior Just for experiment sake I found a work around by running one cycle of training.

...
training_pipeline = (
            dt.Value(crop)
            >> dt.Multiply(lambda: np.random.uniform(0.9, 1.1))
            >> dt.Add(lambda: np.random.uniform(-0.1, 0.1))
            >> dt.MoveAxis(-1, 0)
            >> dt.pytorch.ToTensor(dtype=torch.float32))

training_dataset = dt.pytorch.Dataset(training_pipeline, length=400, replace=False)

dataloader = dl.DataLoader(training_dataset, batch_size=16, shuffle=True)

loaded = dl.LodeSTAR.load_from_checkpoint(self.save_path ,n_transforms=4, optimizer=dl.Adam(lr=1e-4))
loaded.eval()

loaded = dl.LodeSTAR(model=loaded.model,n_transforms=4, optimizer=dl.Adam(lr=1e-4)).build()
trainer = dl.Trainer(max_epochs=1)

trainer.fit(loaded, dataloader)
...

This results in no translation problem but I guess it's rather hacky. Maybe fitting the model saves some transformations that are defined when creating the training dataset?

(sorry for the close, i missclicked)

[BenjaminMidtvedt]

@naikaida That is strange!

Have you also tried loaded.train()? When calling trainer.fit, lightning does put the model in train mode. So that would explain it. If this solves the issue, it should be considered a bug.

[naiikaa]

@BenjaminMidtvedt do you mean it like this?

loaded = dl.LodeSTAR.load_from_checkpoint(<ckpt-path>
,n_transforms=4, optimizer=dl.Adam(lr=1e-4))
loaded.eval()
loaded.train()

I also tried it with leaving the .eval() just calling .train(). In all instances the same problem with the translation of the detections.

[BenjaminMidtvedt]

@naikaida Thank you! I'll give it a shot myself and see if I can replicate it

[BenjaminMidtvedt]

@naikaida I'm unable to reproduce the issue. Simply calling loaded.eval() is enough to make it work for me.

It should be an offset by exactly half the image size. So you can manually correct it for now by doing

x = detections[:, 1] + image.shape[1] / 2
y = detections[:, 0] + image.shape[0] / 2