ClassificationTask predict step

[isaaccorley]

Running trainer.predict() with the ClassificationTask fails because the default pytorch_lightning predict_step expects a batch tuple and not a dict. This PR overrides predict_step to work with our classification datasets.

E.g.

preds = trainer.predict(model=task, dataloaders=datamodule.test_dataloader())

# Or if your datamodule has a `predict_dataloader` method defined
preds = trainer.predict(model=task, datamodule=datamodule)

[calebrob6]

I don't understand the failing test here.

[adamjstewart]

Should we add this to all trainers?

[adamjstewart]

No idea what's going on with the failing test, the latest passing commit on master has the same versions of mypy and numpy, so I don't know why it would be failing all of a sudden.

[adamjstewart]

Technically this could go in 0.3.2 since it doesn't actually add a new feature, it overrides an existing feature.

[isaaccorley]

Looks like it's a bug introduced with the latest mypy that breaks with Python 3.10.7. The fix will come out in the next mypy release. python/mypy#13627

[isaaccorley]

The plan is to add this to all the trainers but each one will have it's own specific needs, e.g you may not want to accumulate predicted segmentation masks into a list in memory and would rather save them in some other way.

[adamjstewart]

That is unfortunate, I guess dependabot can't control exactly which Python version we test. We could pin to 3.10.6 if we want to get the tests working.

[isaaccorley]

@adamjstewart mypy v0.981 is out which includes a fix for this failing test. Should I update requirements/tests.txt in this PR or a separate one?

[adamjstewart]

Dependabot will update it for you if you wait < 24 hrs (I think it checks once per day).

[isaaccorley]

Not sure why codecov is complaining. When I go to the link it only shows coverage dips on scripts I didn't modify.

Closing and reopening did the trick

[adamjstewart]

Why is the predict step so different from the train/val/test steps? Shouldn't it look something like:

Suggested change

	y_hat: Tensor = self(x).softmax(dim=-1)
	return y_hat
	y_hat = self.forward(x)
	y_hat_hard = y_hat.argmax(dim=1)
	return y_hat_hard

[calebrob6]

Because you will frequently want the per-class probabilities instead of just the top class

[adamjstewart]

And that's not true for train/val/test?

[isaaccorley]

The train/val/test steps aren't meant to return data to a user. The operations that happen are only meant to perform the forward pass and compute metrics (which is why the argmax op is needed). The predict step needs to return the raw softmax outputs because we don't want to assume what a user wants to do with the output.

[adamjstewart]

Do we need yhat: Tensor or can we just use yhat?

[calebrob6]

• The recommended way to run inference is with self(x) not self.forward(x) as self(x) calls hooks while self.forward(x) silently doesn't.
• y_hat is a Tensor so I don't see the problem. Breaking something into multiple lines to get mypy to stop complaining is ridiculous overhead.

[adamjstewart]

Should we change other methods to replace self.forward(x) with self(x)? Just want to make sure we're consistent.

[isaaccorley]

Yes we should change. We shouldn't be calling forward anywhere.

[adamjstewart]

P.S. I feel like the yhat: Tensor thing is due to missing type hints for softmax but I couldn't figure out exactly where I need to add those. PyTorch's source code is a maze.

[isaaccorley]

I think I was able to fix it by calling it as a function e.g. torch.softmax(tensor) instead of calling it as a method e.g. tensor.softmax()