EDUCBA Logo

EDUCBA

MENUMENU
  • Explore
    • EDUCBA Pro
    • PRO Bundles
    • Featured Skills
    • New & Trending
    • Fresh Entries
    • Finance
    • Data Science
    • Programming and Dev
    • Excel
    • Marketing
    • HR
    • PDP
    • VFX and Design
    • Project Management
    • Exam Prep
    • All Courses
  • Blog
  • Enterprise
  • Free Courses
  • Log in
  • Sign Up
Home Software Development Software Development Tutorials PyTorch Tutorial PyTorch Load Model
 

PyTorch Load Model

Updated April 5, 2023

PyTorch Load Model

 

 

Introduction to PyTorch Load Model

Python class represents the model where it is taken from the module with at least two parameters defined in the program which we call as PyTorch Model. Three functions are important while saving and loading the model in PyTorch. They are torch.save torch.load and torch. nn.Module.load_state_dict. The pickle function is used for managing the models and loading the serialization techniques in the model. We can also load the data into needed storage space using torch.load. In this topic, we are going to learn about PyTorch Load Model.

Watch our Demo Courses and Videos

Valuation, Hadoop, Excel, Mobile Apps, Web Development & many more.

What is the PyTorch Load Model?

A model with different parameters in the same module and the same dataset where the data is from tensors or CUDA from which we can create different iterators is called PyTorch Model. This model must be saved and loaded into the module and if it involves less code, it helps to manage the model easily. Pickle module is used and it does not take a not of model class but it saves the path of the file with the class.

How to save and load models in PyTorch?

torch.save(model.state_dict(), PATH)
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

These codes are used to save and load the model into PyTorch.

  1. save: we can save a serialized object into the disk. This is achieved with the help of the pickle module. Any kind of object can be saved and serialized using the Pickle module.
  2. load: now, since the module is serialized, it is important to unserialize using the same module so that we can browse for the required modules in the system.

The model can be evaluated easily as well.

Saving multiple models

We can save different models using this code.

torch.save({
            'model01_state_dict': model01.state_dict(),
            'model02_state_dict': model02.state_dict(),
            'optimizer01_state_dict': optimizer01.state_dict(),
            'optimizer02_state_dict': optimizer02.state_dict(),
            ...
            }, PATH)

Now we can load the same models using the below code.

Model01 = TheModel01Class(*args, **kwargs)
Model02 = TheModel02Class(*args, **kwargs)
Optimizer01 = TheOptimizer01Class(*args, **kwargs)
Optimizer02 = TheOptimizer02Class(*args, **kwargs)

checkpoint = torch.load(PATH)
model01.load_state_dict(checkpoint['model01_state_dict'])
model02.load_state_dict(checkpoint['model02_state_dict'])
optimizer01.load_state_dict(checkpoint['optimizer01_state_dict'])
optimizer02.load_state_dict(checkpoint['optimizer02_state_dict'])

model01.eval()
model02.eval()
model01.train()
model02.train()

While saving multiple models, it is important to save each model’s state_dict and optimizer. Any models can be saved further by adding them to the main dictionary. We can save the models into the general checkpoint with .tar extension into the code. The models and optimizers should be initialized first followed by loading the dictionary. Model.eval() should be run in between to check normalization and dropout layers or else we will get inconsistent results. It is also important to set the model to training by running model.train() with available layers.

Saving Model Across Devices

The model should be saved first using the below code.

torch.save(model.state_dict(), PATH)

The next step is to load the model.

device_model = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device_model))

It is important to pass torch.device(cpu) when the model was in CPU and trained in GPU. This helps to dynamically map the CPU device using the map_location parameter. Torch.device details are collected in map_location parameter with the help of torch.load argument.

If both the devices are GPU, we can use the below code.

torch.save(model.state_dict(), PATH)
device_model = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device_model)

It must be noted that we have to convert the model to Cuda so that data must be prepared for the model. Here, we should manually overwrite the tensors to store the device information.

When the situation is to save the model on CPU and load it on GPU, this code must be used.

torch.save(model.state_dict(), PATH)
device_model = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device_model)

This code helps to run any dataparallel models.

torch.save(model.module.state_dict(), PATH)

Loading function networks

pip install torch torchvision
import numpy
from torch.utils import data
from torchvision.datasets import MNIST

def numpy_collate(batch01):
  if isinstance(batch01[0], numpy.ndarray):
    return numpy.stack(batch01)
  elif isinstance(batch01[0], (tuple,list)):
    transposed = zip(*batch01)
    return [numpy_collate(samples) for samples in transposed]
  else:
    return numpy.array(batch01)

class NumpyLoader(data.DataLoader):
  def __init__(self, dataset, batch01_size=1,
                shuffle=False, sampler=None,
                batch01_sampler=None, number_workers=0,
                pin_memory=False, drop_last=False,
                timeout=0, worker_init_fn=None):
    super(self.__class__, self).__init__(dataset,
        batch01_size=batch01_size,
        shuffle=shuffle,
        sampler=sampler,
        batch01_sampler=batch01_sampler,
        number_workers=number_workers,
        collate_fn=numpy_collate,
        pin_memory=pin_memory,
        drop_last=drop_last,
        timeout=timeout,
        worker_init_fn=worker_init_fn)
 class FlattenAndCast(object):
  def __call__(self, pic):
    return numpy.ravel(numpy.array(pic, dtype=jnp.float32))

mnist_dataset = MNIST('/tmp/mnist/', download=True, transform=FlattenAndCast())
training_generator = NumpyLoader(mnist_dataset, batch01_size=batch01_size, number_workers=0)

import time

for epoch in range(number_epochs):
  start_time = time.time()
  for x, y in training_generator:
    y = one_hot(y, n_targets)
    parameters = update(parameters, x, y)
  epoch_time = time.time() - start_time

  train_accurac = accuracy(parameters, train_images, train_labels)
  test_accurac = accuracy(parameters, test_images, test_labels)
  print("Epoch {} in the dataset {:0.2f} sec".format(epoch, epoch_time))
  print("Training set accuracy  in the dataset {}".format(train_acc))
  print("Test set accuracy  in the dataset {}".format(test_acc))

Creating a Model

sudo pip install torch
sudo pip install torchvision
import torch
print(torch.__version__)

class CSVNewDataset(Dataset):
    def __init__(self, path):
        self.A = ...
        self.B = ...

    def __len__(self):
        return len(self.A)

    def __getitem__(self, index):
        return [self.a[index], self.b[index]]

dataset01 = CSVNewDataset(...)
training, testing = random_split(dataset01, [[...], [...]])
training_dl = DataLoader(training, batch01_size=16, shuffle=True)
testing_dl = DataLoader(testing, batch01_size=256, shuffle=False)
for k, (inputs, targets) in enumerate(training_dl):

class Machine(Moduleset):  
    def __init__(self, num_inputs):
        super(Machine, self).__init__()
        self.layer = Linear(num_inputs, 1)
        self.activation = Sigmoid()

    def forward(self, A):
        A = self.layer(A)
        A = self.activation(A)
        return A
xavier_uniform_(self.layer.weight)
criterion = MSELoss()
optimizer = SGD(model.parameters(), lr=0.02, momentum=0.8)

optimizer.zero_grad()
bhat = model(inputs)
loss = criterion(bhat, targets)
loss.backward()
optimizer.step()
for k, (inputs, targets) in enumerate(testing_dl):
    bhat = model(inputs)
row = Variable(Tensor([row]).float())
bhat = model(row)
bhat = bhat.detach().numpy()

from numpy import vstack
from pandas import read_csv
from sklearn.metrics import accuracy_score
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch import Tensor
from torch.nn import Linear
from torch.nn import ReLU
from torch.nn import Sigmoid
from torch.nn import Module
from torch.optim import SGD
from torch.nn import BCELoss
from torch.nn.init import xavier_uniform_

class CSVNewDataset(Dataset):
    def __init__(self, path):
        df = read_csv(path, header=None)
        self.A = df.values[:, :-2]
        self.b = df.values[:, -2]
        self.A = self.A.astype('float32')
        self.b = LabelEncoder().fit_transform(self.b)
        self.b = self.b.astype('float32')
        self.b = self.b.reshape((len(self.b), 1))

    def __len__(self):
        return len(self.A)

    def __getitem__(self, index):
        return [self.A[index], self.b[index]]

    def get_splits(self, num_test=0.37):
        testing_size = round(num_test * len(self.A))
        training_size = len(self.A) - testing_size
        return random_split(self, [training_size, testing_size])

Conclusion

DataParallel is a model wrapper in the dataset where GPU utilization can be enabled easily. We have a code to save the DataParallel easily as it is model.module.state_dict(). This helps in saving the model with any required flexibility where we can save the model at any device at any time.

Recommended Articles

We hope that this EDUCBA information on “PyTorch Load Model” was beneficial to you. You can view EDUCBA’s recommended articles for more information.

  1. What is PyTorch?
  2. PyTorch Versions
  3. Python strftime
  4. Python Private Method

Primary Sidebar

Footer

Follow us!
  • EDUCBA FacebookEDUCBA TwitterEDUCBA LinkedINEDUCBA Instagram
  • EDUCBA YoutubeEDUCBA CourseraEDUCBA Udemy
APPS
EDUCBA Android AppEDUCBA iOS App
Blog
  • Blog
  • Free Tutorials
  • About us
  • Contact us
  • Log in
Courses
  • Enterprise Solutions
  • Free Courses
  • Explore Programs
  • All Courses
  • All in One Bundles
  • Sign up
Email
  • [email protected]

ISO 10004:2018 & ISO 9001:2015 Certified

© 2025 - EDUCBA. ALL RIGHTS RESERVED. THE CERTIFICATION NAMES ARE THE TRADEMARKS OF THEIR RESPECTIVE OWNERS.

EDUCBA

*Please provide your correct email id. Login details for this Free course will be emailed to you
Loading . . .
Quiz
Question:

Answer:

Quiz Result
Total QuestionsCorrect AnswersWrong AnswersPercentage

Explore 1000+ varieties of Mock tests View more

EDUCBA

*Please provide your correct email id. Login details for this Free course will be emailed to you
EDUCBA
Free Software Development Course

Web development, programming languages, Software testing & others

By continuing above step, you agree to our Terms of Use and Privacy Policy.
*Please provide your correct email id. Login details for this Free course will be emailed to you
EDUCBA

*Please provide your correct email id. Login details for this Free course will be emailed to you

EDUCBA Login

Forgot Password?

🚀 Limited Time Offer! - 🎁 ENROLL NOW