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
Or from source
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.2.tar.gz
(7.1 kB
view hashes)
Built Distribution
Close
Hashes for model_loads-0.2.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 34b7ef9a5a2f16c016e03e69278d8734492b214b3c2ee106e8b55c78fd0eb34c |
|
MD5 | a599b523612d0a5d9e183e18bd4b7d0c |
|
BLAKE2b-256 | 023df387f05ad22be5f7e6351759e0b4cd1d5cb27ca511888d3b2ff609298f41 |