load_state_dict does not return the model

[konstantinklemmer]

Fixed an error in the state dict loading of the tutorial

[adamjstewart]

Not sure why tests are failing but it's clearly unrelated to this PR. Will try to investigate in a separate PR.

[adamjstewart]

Fixed an error in the state dict loading of the turorial

What's the error?

[calebrob6]

@konstantinklemmer ping

[adamjstewart]

P.S. I fixed the failing test, tests should pass after updating your branch. I'm just curious why this PR is needed since it doesn't fail during testing.

[adamjstewart]

It's still unclear to me the purpose of this PR. What is the bug it is trying to solve?

[calebrob6]

There isn't a bug, this is just an example of how to load pre-trained models differently.

[adamjstewart]

That's not what the PR descriptions says...

[isaaccorley]

The error that's being fixed is that model.load_state_dict() returns either None or a list of incompatible keys. Our example incorrectly does model = model.load_state_dict which overrides the model itself to None or a list.

[isaaccorley]

This looks okay to me

[adamjstewart]

Thanks @isaaccorley, I understand the error now.

It looks like we actually make this same mistake in the README. And we define our own custom torchgeo.trainers.utils.load_state_dict that behaves differently from the builtin one. We should fix all of these at the same time.

I'm not sure about the num_classes comment. We aren't trying to teach people how to use timm, just how to use TorchGeo. I don't disagree that it's useful, just that it's in the wrong place.

@konstantinklemmer let me know if you want me to make these changes myself. If I don't hear back I'll assume this PR has been abandoned and take over.

[konstantinklemmer]

Sorry, no idea why I am not getting notifications for this PR...

Yes, I am happy to remove that comment and the example with num_classes=0. Does that sound alright? Basically just change model = model.load_state_dict() to model.load_state_dict().

[adamjstewart]

Yes, and update the README and our builtin load_state_dict as well. Let me know if you want help with the latter.

[konstantinklemmer]

Yes please, I am not sure how to tackle the builtin load_state_dict problem (or what it even is exactly).

[adamjstewart]

or what it even is exactly

It's not really a problem per se, just that we define a wrapper around load_state_dict that returns the model instead of returning hits/misses like the builtin PyTorch method. It's kind of confusing to have different behavior, so I think we should match the behavior of the builtin method. Luckily, this method isn't public facing, so this change won't be backwards incompatible.

[konstantinklemmer]

Okay! So there are two functions that I found that, if I understand correctly, are relevant:

• torchgeo/torchgeo/trainers/utils.py

  Line 74 in 436baa9

  def load_state_dict(model: Module, state_dict: "OrderedDict[str, Tensor]") -> Module:

• torchgeo/tests/trainers/test_utils.py

  Line 42 in 436baa9

  def test_load_state_dict(checkpoint: str, model: Module) -> None:

  
  Is that correct?

In both cases, if we want to keep them as standalone functions that return a model with weights loaded, we should probably rename them?

[adamjstewart]

torchgeo.trainers.utils.load_state_dict is the one you would change. The test_load_state_dict is just to test the method, the inputs don't matter.

We can either A) change the name, or B) change the return value to match torch.nn.Module.load_state_dict. If the latter works, I would actually prefer the latter.

[konstantinklemmer]

Well, load_state_dict does not have a return value, no? So for B) we'd just need to get rid of the return call, i.e. this line here (unless I misunderstand): https://github.com/microsoft/torchgeo/blob/436baa9773977c789152854ac7b4eff90e0d9e95/torchgeo/trainers/utils.py#L119C5-L119C17

Then load_state_dict(model, ckpt['state_dict']) should be equivalent to model.load_state_dict(ckpt['state_dict']).

[adamjstewart]

The builtin load_state_dict returns missing and unexpected keys: https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict

So you'll just return the output of that call so that our wrapper matches.

[konstantinklemmer]

Okay after some more digging, the nn.Module.load_state_dict returns an object _IncompatibleKeys (https://github.com/pytorch/pytorch/blob/c3780010a58a84920335296ee5f091a0db18259f/torch/nn/modules/module.py#L29). This object uses missing and incompatible keys. However, when you check the load_state_dict and _load_from_state_dict (https://github.com/pytorch/pytorch/blob/c3780010a58a84920335296ee5f091a0db18259f/torch/nn/modules/module.py#L1953) functions, it seems that both will be empty dicts:

missing_keys: List[str] = []
unexpected_keys: List[str] = []

unless the strict=True flag is on. The torchgeo.trainers.utils.load_state_dict does not use that flag. So should I just return _IncompatibleKeys with those empty dicts?

[adamjstewart]

Let's just return the output of nn.Module.load_state_dict. For type hints, try:

-> Tuple[List[str], List[str]]

I think that will work, and be correct. The builtin method has no return type annotations so it shouldn't complain that we aren't using a NamedTuple. If you have any trouble with mypy, let me know and I can hack on it.

[konstantinklemmer]

Got it! Will try if that works and report back.

[konstantinklemmer]

Ok I think I updated all files (the README, the notebook, the utils.py) but tests are failing.

[konstantinklemmer]

Yay! Thanks for helping with this - seems though that the original problem with the notebook test still persists?

[adamjstewart]

Completely unrelated problem, fixed in #1838