Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.80%. Comparing base (c4dd472) to head (0c5b02e).

@@            Coverage Diff             @@
##             main    #1915      +/-   ##
==========================================
- Coverage   96.81%   96.80%   -0.01%     
==========================================
  Files         172      172              
  Lines        8444     8515      +71     
==========================================
+ Hits         8175     8243      +68     
- Misses        269      272       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@SiddhantBahuguna maybe from_weights a better method name or something different ? ^^

@felixdittrich92 Since most of the implementations are the exact same, wouldn't it be better to just have a common definition & just import it when needed (ofc keep the different ones in their respective classes)? I feel like having this many duplicated bits is a sure way to forget about it when it matters down the road 😅

Edit: misread the implementation. But it does feel confusing that from_pretrained implementations are all made into methods just to relay params to load_pretrained_params? Why not just add the non-standard implementations to load_pretrained_params instead?

@felixdittrich92 Since most of the implementations are the exact same, wouldn't it be better to just have a common definition & just import it when needed (ofc keep the different ones in their respective classes)? I feel like having this many duplicated bits is a sure way to forget about it when it matters down the road 😅

Edit: misread the implementation. But it does feel confusing that from_pretrained implementations are all made into methods just to relay params to load_pretrained_params? Why not just add the non-standard implementations to load_pretrained_params instead?

@sebastianMindee
Yeah in general the cleanest way would be to have a base class which implements already the functionality and all models inerhit from them but this would require a larger refactoring as mentioned and I think would raise an circular import issue.

So the idea of the wrapper method was that we can use this place for custom model logic (like parseq in this case).
But I agree we could also modify load_pretrained_params slighty make it private, import from models.utils and bind it afterwards model.load_pretrained_params = types.MethodTypes(model, '_load_pretrained_params') in this case we would reduce the most boilerplate code ..I tried to avoid this case as much as possible because anyway it feels a bit hacky so I used it only at places where a larger refactoring would be required 😅
wdyt ?

btw we wouldn't miss to implement this method because if we add the corresponding test entries the CI would fail (Added a unittest to check that all models have the from_pretrained method)

@sebastianMindee

After double checking I have to revert the above case because it would result with the same issue as modifying load_pretrained_params directly - results in circular import

And I'm not open to switch to full lazy loading by skipping anything non matching 😅

from ..recognition import PARSeq  # circular

def load_pretrained_params(
    model: nn.Module,
    path_or_url: str | None = None,
    hash_prefix: str | None = None,
    ignore_keys: list[str] | None = None,
    **kwargs: Any,
) -> None:
    """Load a set of parameters onto a model

    >>> from doctr.models import load_pretrained_params
    >>> load_pretrained_params(model, "https://yoursource.com/yourcheckpoint-yourhash.pt")

    Args:
        model: the PyTorch model to be loaded
        path_or_url: the path or URL to the model parameters (checkpoint)
        hash_prefix: first characters of SHA256 expected hash
        ignore_keys: list of weights to be ignored from the state_dict
        **kwargs: additional arguments to be passed to `doctr.utils.data.download_from_url`
    """
    if path_or_url is None:
        logging.warning("No model URL or Path provided, using default initialization.")
        return

    archive_path = (
        download_from_url(path_or_url, hash_prefix=hash_prefix, cache_subdir="models", **kwargs)
        if validators.url(path_or_url)
        else path_or_url
    )

    if isinstance(model, PARSeq):
        # NOTE: This is required to make the model backward compatible with already trained models docTR version <0.11.1
        # ref.: https://github.com/mindee/doctr/issues/1911
        if ignore_keys is None:
            ignore_keys = []
        ignore_keys.extend([
            "decoder.attention_norm.weight",
            "decoder.attention_norm.bias",
            "decoder.cross_attention_norm.weight",
            "decoder.cross_attention_norm.bias",
        ])

...

An option would be to check only for the name 🙈

if type(model).__name__ == "PARSeq":

Ok finally:

Option 1:

model = vitstr()
model.from_pretrained("path/to/weights.X")

We keep it as provided in the PR

Advantage:

• Clear structure for each model & class bound logic (no condition required about model type / name)
• Already known interface from a lot of other libs like transformers / sentence-transformers and so on -> class method .from_pretrained

Disadvantage:

• Produces some boilerplate code & a wrapper method
• Uses method binding at places where a ("to class wrapper") refactoring would be required otherwise

Option 2:

We handle anything in load_pretrained_params

on user side then it would be:

model = vitstr()
load_pretrained_params(model, "path/to/weights.X")

Advantage:

• Avoids all the boilerplate
• Minimal changes required

Disadvantage:

• Model name condition bound in load_pretrained_params
• Feels less familiar to well known libs
• Provides additional options which could be a bit confusing for a simple loading function (hash_prefix, ignore_keys, ..) which are "masked" as kwargs in option 1

@sebastianMindee @SiddhantBahuguna Now we should decide which option would you prefer ? 😄

My vote goes to option 1 because from user view it feels a bit smarter 😅 AND in this case we have well unittested that each model provides a defined method to load weights - this case provides several checks - pretrained=True would fail if the model does not implement the method (because of the inner model.from_pretrained call if pretrained=True is passed) and the added assert hasattr checks also

LGTM but I'd like to wait for @SiddhantBahuguna's approval before giving it the official go 🙂

Yeah that's fine ☺️

Will update the missing part for #1912 next week :)

Done :)