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:
pip install model-loads
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))
Done:
load model which save model by
torch.save(model, "path/to/model")
Todo list
DATA PARALLELISM
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.2.0.tar.gz
(7.0 kB
view hashes)
Built Distribution
Close
Hashes for model_loads-0.2.0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 7d4452a12e32181f1528db2bea4152d538c0bca0fe3176d28bdb34d9769e780e |
|
MD5 | a326df5aec44505d94ea78a383c3a535 |
|
BLAKE2b-256 | 934cd3157445a27b5b582b71a244ef718a03ac8a8cf37189b4336711185ae3a7 |