Loads GPU or CPU pytorch models
Project description
model_loads is an open-source Python package for pytorch load models easy.
PyTorch is a Python package that provides two high-level features:
- Tensor computation (like NumPy) with strong GPU acceleration
- Deep neural networks built on a tape-based autograd system
It's annoying to load cpu model to gpu devices or load multi-gpus trained model to single gpu devices sometimes, And this package try to simplify it.
Table of Contents
Installation
To install load_models, you can do as follow:
git clone https://github.com/cwh94/model_loads.git
cd load_models
python setup.py bdist_egg
python setup.py install
Getting Started
- load pth model to GPU device
import model_loads as lo
import torchvision.models as models
model = models.MobileNetV2()
model_path = "../examples/models/pth/mobilenet_v2-b0353104.pth"
lo.load_models(model_path, model, use_gpu=True)
print(model)
print(type(model))
- load tar model(which contains state_dict and optimization info or accuracy) to CPU device
from models.tar.mobilenet_v2 import MobileNetV2
model = MobileNetV2()
model_path = "models/tar/checkpoint.pth.tar"
lo.load_models(model_path, model)
print(model)
- load model to CPU device
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
model = models.MobileNetV2()
model_path = "models/pth/mobilenet_v2-b0353104.pth"
lo.load_models(model_path, model, use_gpu=True)
print(model)
print(type(model))
Todo list
load model which save model by
torch.save(model, "path/to/model")
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
model-loads-0.1.1.tar.gz
(3.4 kB
view hashes)
Built Distribution
Close
Hashes for model_loads-0.1.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3053900154133b2a5387f0da284085778df3a465075337c95a8d97de924224aa |
|
MD5 | aa036c85a3f4605fbe67485bef579920 |
|
BLAKE2b-256 | c468f891a655bc12c90edb28d071a6f57ecc40c93504566389e8bd26a07cb9a8 |