Kitware / SMQTK

Python toolkit for pluggable algorithms and data structures for multimedia-based machine learning.
http://smqtk.readthedocs.org/
Other
75 stars 58 forks source link

Auto fetch pytorch model and return layer #408

Closed vbhavank closed 4 years ago

vbhavank commented 4 years ago

Using PyTorch model currently requires the user to define the architecture for some commonly used models. Therefore, we make use of torchvision library to import PyTorch models using labels specifying the name of the model (i.e. resnet50, vgg16) and return layer for feature collection (i.e. avgpool, classifier.4).