Skip to main content

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

  1. 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))
  1. 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)
  1. 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.

Files for model-loads, version 0.1.1
Filename, size File type Python version Upload date Hashes
Filename, size model_loads-0.1.1-py3-none-any.whl (9.1 kB) File type Wheel Python version py3 Upload date Hashes View hashes
Filename, size model-loads-0.1.1.tar.gz (3.4 kB) File type Source Python version None Upload date Hashes View hashes

Supported by

Elastic Elastic Search Pingdom Pingdom Monitoring Google Google BigQuery Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page