11"""
22Saving and loading multiple models in one file using PyTorch
33============================================================
4-
54Saving and loading multiple models can be helpful for reusing models
65that you have previously trained.
76
87Introduction
98------------
10-
119When saving a model comprised of multiple ``torch.nn.Modules``, such as
1210a GAN, a sequence-to-sequence model, or an ensemble of models, you must
1311save a dictionary of each model’s state_dict and corresponding
1412optimizer. You can also save any other items that may aid you in
1513resuming training by simply appending them to the dictionary.
16-
1714To load the models, first initialize the models and optimizers, then
1815load the dictionary locally using ``torch.load()``. From here, you can
1916easily access the saved items by simply querying the dictionary as you
2017would expect.
21-
2218In this recipe, we will demonstrate how to save multiple models to one
2319file using PyTorch.
2420
2521Setup
2622-----
27-
2823Before we begin, we need to install ``torch`` if it isn’t already
2924available.
3025
26+ ::
27+
28+ pip install torch
29+
3130"""
3231
33- pip install torch
3432
3533
3634######################################################################
4341# 4. Save multiple models
4442# 5. Load multiple models
4543#
46- # **1) Import necessary libraries for loading our data**
44+ # 1. Import necessary libraries for loading our data
4745# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
4846#
4947# For this recipe, we will use ``torch`` and its subsidiaries ``torch.nn``
5654
5755
5856######################################################################
59- # **2) Define and intialize the neural network**
57+ # 2. Define and intialize the neural network
6058# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
6159#
6260# For sake of example, we will create a neural network for training
@@ -88,7 +86,7 @@ def forward(self, x):
8886
8987
9088######################################################################
91- # **3) Initialize the optimizer**
89+ # 3. Initialize the optimizer
9290# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
9391#
9492# We will use SGD with momentum to build an optimizer for each model we
@@ -100,7 +98,7 @@ def forward(self, x):
10098
10199
102100######################################################################
103- # **4) Save multiple models**
101+ # 4. Save multiple models
104102# ~~~~~~~~~~~~~~~~~~~~~~~~~~~
105103#
106104# Collect all relevant information and build your dictionary.
@@ -118,7 +116,7 @@ def forward(self, x):
118116
119117
120118######################################################################
121- # **5) Load multiple models**
119+ # 4. Load multiple models
122120# ~~~~~~~~~~~~~~~~~~~~~~~~~~~
123121#
124122# Remember to first initialize the models and optimizers, then load the
@@ -161,4 +159,4 @@ def forward(self, x):
161159#
162160# - TBD
163161# - TBD
164- #
162+ #
0 commit comments