[Feature Proposal] Extend torch hub to better support cloud serving and edge deployment

[RalphMao]

🚀 The feature, motivation and pitch

TL;DR: Extend torch hub to better support cloud serving and edge deployment by:

1. Extend the hubconf.py as an entrypoint for serving and deployment, like HuggingFace Inferece Handler
2. (Optional) Allow third-party infra providers to integrate model deployment and performance reporting to the hub website.

Background

I was a PhD researcher a while ago and am now an AI practitioner in the industry. From my personal experience and observation, the pytorch hub provides incremental value to researchers, because:

1. Researchers usually rely on conferences and their own networks to find the latest models.
2. Researchers are familiar with cloning a repo, peeking into the source and doing modifications. The simple torch.hub.load interface is not very much helpful to them.

However, industry community usually don't have the advantages mentioned above, therefore the hub can be a valuable information source to them. But when they choose a model, they not only want an interface to run some sample data, deployment easiness and model performance are also important factors to consider.

That's why huggingface transformers, yolo-series, and mmdetection are very popular.
But not all researchers like to submit their models to those repos due to the extensive customizations needed, therefore, a lightweight hubconf.py could just reduce the gap between research and production need.

Proposed changes

hubconf.py: provides a reference format to share models, and acts as an interface for deployment, benchmarking and serving.

A very rough example is shown below, just for discussion purpose.

__model__ = ["MyModel"]

def MyModel:
  def __init__(self, pretrained=False, **kwargs): 
    self.model = build_model(pretrained=pretrained)
  
  def forward(self, x: Union[Tuple[Any, ...], torch.Tensor]) -> Union[Tuple[Any, ...], torch.Tensor]:  
    # Core tensor-in, tensor-out model function
    return self.model(x)
  
  def get_input(self): 
    # get sample data
    return preprocess(cv2.imread('assets/img1.png'))

  def inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
    # optional, for model serving only
    output = self.model(preprocess(data['input']))
    return {"predictions": postprocess(output)}

Hub website: "Discover and publish models to a pre-trained model repository designed for research exploration." -> "For research exploration and production use". And in the future, allow third-party providers to plug in model statistics based on hubconf.py

Disclaimer
Our company is currently working on an MLOps platform to facilitate the production of PyTorch models. We see a clear benefit that a shared protocol could further reduce the friction between research and industry communities. We would like to hear feedback from PyTorch team and we are also open to take the initial effort to extend the hubconf and adopt open-source repos to fit this extended format.

Alternatives

No response

Additional context

No response

cc @nairbv @NicolasHug @vmoens @jdsgomes

[vmoens]

Hi @RalphMao, sorry for the delay in responding to you.

Thanks for posting this. I'm trying to understand precisely what you're proposing, feel free to correct me if I'm making wrong assumptions.
Are you suggesting that every model registered in torch.hub should be made compliant with this structure, or would it be optional?
In the former case it would be a bit difficult to bring in, as it'd require us to go to multiple contributors and ask them to change their code accordingly. I would not bet on many of them doing it.

Also, we provide a hub for models provided by third parties without forcing them to give a pretrained version (although many do provide this). In this context it's not obvious what a benchmark execution would mean (since the model is not tied to a specific dataset or task). Same for self.get_input: what would this mean without a precise context?

Let me know if I got anything wrong, I'm not sure I clearly understood the kind of changes you were suggesting.