[WIP] Integration with DeepLabCut 3.0 - PyTorch Engine

[n-poulsen]

This pull requests updates DeepLabCut-Live for models exported with DeepLabCut 3.0. TensorFlow models can still be used, and the code is siloed so that only the engine used to run the code is required as a package (i.e. no need to install TensorFlow if you want to run live pose estimation with PyTorch models).

If you want to give this PR a try, you can install the code in your local conda environment by running:

pip install "git+https://github.com/DeepLabCut/DeepLabCut-live.git@dlclive3"

[maximpavliv]

Suggested change

	- Feb 2021: DeepLabCut-Live! was featured in **Nature Methods**: [
	"Real-time behavioral analysis"](https://www.nature.com/articles/s41592-021-01072-z)
	- Feb 2021: DeepLabCut-Live! was featured in **Nature Methods**:
	["Real-time behavioral analysis"](https://www.nature.com/articles/s41592-021-01072-z)

[maximpavliv]

DLCLive's constructor doesn't have a max_detections argument

[maximpavliv]

Shouldn't we check if has_tqdm here?

[maximpavliv]

Let's name it vid_writer as in benchmark()? I find vid_writer indeed better

[sneakers-the-rat]

am only a peripheral maintainer here at this point, but i see that we have copied and pasted this whole directory? https://github.com/DeepLabCut/DeepLabCut/tree/main/deeplabcut/pose_estimation_pytorch

that sounds pretty fragile to me! if we're already copy/pasting it once, we might as well just make that a cut/paste into a separate deeplabcut-pytorch package that both deeplabcut and deeplabcut-live can depend on, right?

[MMathisLab]

Hey @sneakers-the-rat afaik its not strictly the same though; this the same reason we don't keep dlc-live in DeepLabCut main. We also used to have core, but it becomes yet another package to maintain and ultimately was more work for dev team. But I would leave it to @n-poulsen to give a better technical answer (as this addition here I haven't been deeply involved in yet) ❤️- thanks for your engagement!!!

[sneakers-the-rat]

oh ya i don't see a reason to have dlc-live in main dlc, i would think that the dep graph would look something like this

flowchart TB
    deeplabcut --> dlc-pytorch
    deeplabcut --> dlc-tensorflow
    dlc-pytorch --> dlclibrary
    dlc-tensorflow --> dlclibrary
    dlc-live --> dlc-pytorch
    dlc-live --> dlc-tensorflow
    dlc-live --> dlclibrary
    deeplabcut --> dlclibrary

Loading

where dlclibrary has stuff that is boilerplate or super lightweight and common to all the packages (like exceptions, utility functions, and so on), and then the dlc-pytorch and dlc-tensorflow have all the framework-specific stuff and models in them - this would also make a single-source of dependency compatibility ranges which i understand to be like an infinite problem. and then dlc-live and deeplabcut are the application layer that wraps the raw models and etc. up into user-facing tools.

We also used to have core, but it becomes yet another package to maintain and ultimately was more work for dev team

I believe y'all have your own processes and flow and etc. and don't mean to step on that. the thing that works is the thing that works. i would suspect that having multiple copies of what was initially the same code drift apart (or require active duplication as changes are made) would be more maintenance labor in all but the short term, but i also don't deal with tf or torch directly and am aware they introduce their own packaging nightmares wherever they go, so up to y'all ofc.

The thing that we did for the tensorflow stuff is just operate on serialized versions of the models, so the model code for tf isn't in the package at all - does torch not have a similar system for serializing and deserializing models? that worked out really well, and i partly credit that design decision with why this package has needed relatively little maintenance over its lifespan. that is a third way that would avoid new package but also not require a copy/paste :)

[n-poulsen]

Hi @sneakers-the-rat! Thanks for the input! When it comes to duplicating code, only the models from the deeplabcut/pose_estimation_pytorch folder were copied over. This is absolutely a sub-optimal solution, and requires maintenance when new models are added to DeepLabCut or bug fixes are made. This was done to have a prototype ready, but I'm still exploring better solutions. The issue is that:

• PyTorch recommends exporting the state_dict for the model, which requires the model class to run inference
• Serialization does exist with PyTorch through TorchScript, but
  1. TorchScript is no longer in active development
  2. It will be replaced by torch.export but it is a prototype under active development and there WILL BE BREAKING CHANGES in the future. I've seen they want to use it to replace TorchScript in 2025, but this makes it harder for me to want to use it. It's also currently not well integrated for torchvision models, so object detectors could not be exported. This will likely lead to issues with other operations we have in our models.
  3. TorchScript can have the same issue with exporting models (when some non-standard operations are used). I'd need to test all of our model architectures to see which ones break - this will just take a bit of time but should be feasible
  4. Not all pose models implemented for the PyTorch engine output heatmaps: so to go from the model output to predicted pose, we may require different predictors. I want to experiment with exporting predictor functions alonside the models (so they are fully serializable, taking an image as input and outputting an array of shape (num_individuals, num_keypoints, 3) containing the predicted pose) but that may take some time as well. An example of a model returning a different type of output would be RTMPose, which uses "coordinate classification".
• Another option would be to serialize through ONNX, but the same issue with the predictors remain. I'm exploring this as well.

If we can't find a good way to serialize models, it's true that the best solution may be to have a separate package containing the models and predictors - but that's something I still want to explore.

If you have any opinions on the best way to this I'm glad to get any other input!

[sneakers-the-rat]

aha, that is extremely annoying there is no good way to serialize pytorch models, or at least serialization is in a transitional period right now.

I'm not gonna be maintaining it, and have said my piece re: splitting off into a separate package, so by all means! good to see this update getting ported over to dlc live, glhf with the rest of the impl, lmk if a review would be helpful <3