Skip to content

Commit acdc91b

Browse files
svekarsc-p-i-o
authored andcommitted
Add weights_only=True to torch.load (#3012)
* Add weights_only=True to torch.load
1 parent 9d97d8f commit acdc91b

22 files changed

+50
-38
lines changed

advanced_source/dynamic_quantization_tutorial.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,8 @@ def tokenize(self, path):
151151
model.load_state_dict(
152152
torch.load(
153153
model_data_filepath + 'word_language_model_quantize.pth',
154-
map_location=torch.device('cpu')
154+
map_location=torch.device('cpu'),
155+
weights_only=True
155156
)
156157
)
157158

advanced_source/static_quantization_tutorial.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ We next define several helper functions to help with model evaluation. These mos
286286
287287
def load_model(model_file):
288288
model = MobileNetV2()
289-
state_dict = torch.load(model_file)
289+
state_dict = torch.load(model_file, weights_only=True)
290290
model.load_state_dict(state_dict)
291291
model.to('cpu')
292292
return model

beginner_source/basics/quickstart_tutorial.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,7 @@ def test(dataloader, model, loss_fn):
216216
# the state dictionary into it.
217217

218218
model = NeuralNetwork().to(device)
219-
model.load_state_dict(torch.load("model.pth"))
219+
model.load_state_dict(torch.load("model.pth", weights_only=True))
220220

221221
#############################################################
222222
# This model can now be used to make predictions.

beginner_source/basics/saveloadrun_tutorial.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,14 @@
3232
##########################
3333
# To load model weights, you need to create an instance of the same model first, and then load the parameters
3434
# using ``load_state_dict()`` method.
35+
#
36+
# In the code below, we set ``weights_only=True`` to limit the
37+
# functions executed during unpickling to only those necessary for
38+
# loading weights. Using ``weights_only=True`` is considered
39+
# a best practice when loading weights.
3540

3641
model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
37-
model.load_state_dict(torch.load('model_weights.pth'))
42+
model.load_state_dict(torch.load('model_weights.pth', weights_only=True))
3843
model.eval()
3944

4045
###########################
@@ -50,9 +55,14 @@
5055
torch.save(model, 'model.pth')
5156

5257
########################
53-
# We can then load the model like this:
58+
# We can then load the model as demonstrated below.
59+
#
60+
# As described in `Saving and loading torch.nn.Modules <pytorch.org/docs/main/notes/serialization.html#saving-and-loading-torch-nn-modules>`__,
61+
# saving ``state_dict``s is considered the best practice. However,
62+
# below we use ``weights_only=False`` because this involves loading the
63+
# model, which is a legacy use case for ``torch.save``.
5464

55-
model = torch.load('model.pth')
65+
model = torch.load('model.pth', weights_only=False),
5666

5767
########################
5868
# .. note:: This approach uses Python `pickle <https://docs.python.org/3/library/pickle.html>`_ module when serializing the model, thus it relies on the actual class definition to be available when loading the model.

beginner_source/blitz/cifar10_tutorial.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ def forward(self, x):
221221
# wasn't necessary here, we only did it to illustrate how to do so):
222222

223223
net = Net()
224-
net.load_state_dict(torch.load(PATH))
224+
net.load_state_dict(torch.load(PATH, weights_only=True))
225225

226226
########################################################################
227227
# Okay, now let us see what the neural network thinks these examples above are:

beginner_source/fgsm_tutorial.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ def forward(self, x):
192192
model = Net().to(device)
193193

194194
# Load the pretrained model
195-
model.load_state_dict(torch.load(pretrained_model, map_location=device))
195+
model.load_state_dict(torch.load(pretrained_model, map_location=device, weights_only=True))
196196

197197
# Set the model in evaluation mode. In this case this is for the Dropout layers
198198
model.eval()

beginner_source/saving_loading_models.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@
153153
# .. code:: python
154154
#
155155
# model = TheModelClass(*args, **kwargs)
156-
# model.load_state_dict(torch.load(PATH))
156+
# model.load_state_dict(torch.load(PATH), weights_only=True)
157157
# model.eval()
158158
#
159159
# .. note::
@@ -206,7 +206,7 @@
206206
# .. code:: python
207207
#
208208
# # Model class must be defined somewhere
209-
# model = torch.load(PATH)
209+
# model = torch.load(PATH, weights_only=False)
210210
# model.eval()
211211
#
212212
# This save/load process uses the most intuitive syntax and involves the
@@ -290,7 +290,7 @@
290290
# model = TheModelClass(*args, **kwargs)
291291
# optimizer = TheOptimizerClass(*args, **kwargs)
292292
#
293-
# checkpoint = torch.load(PATH)
293+
# checkpoint = torch.load(PATH, weights_only=True)
294294
# model.load_state_dict(checkpoint['model_state_dict'])
295295
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
296296
# epoch = checkpoint['epoch']
@@ -354,7 +354,7 @@
354354
# optimizerA = TheOptimizerAClass(*args, **kwargs)
355355
# optimizerB = TheOptimizerBClass(*args, **kwargs)
356356
#
357-
# checkpoint = torch.load(PATH)
357+
# checkpoint = torch.load(PATH, weights_only=True)
358358
# modelA.load_state_dict(checkpoint['modelA_state_dict'])
359359
# modelB.load_state_dict(checkpoint['modelB_state_dict'])
360360
# optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
@@ -407,7 +407,7 @@
407407
# .. code:: python
408408
#
409409
# modelB = TheModelBClass(*args, **kwargs)
410-
# modelB.load_state_dict(torch.load(PATH), strict=False)
410+
# modelB.load_state_dict(torch.load(PATH), strict=False, weights_only=True)
411411
#
412412
# Partially loading a model or loading a partial model are common
413413
# scenarios when transfer learning or training a new complex model.
@@ -446,7 +446,7 @@
446446
#
447447
# device = torch.device('cpu')
448448
# model = TheModelClass(*args, **kwargs)
449-
# model.load_state_dict(torch.load(PATH, map_location=device))
449+
# model.load_state_dict(torch.load(PATH, map_location=device, weights_only=True))
450450
#
451451
# When loading a model on a CPU that was trained with a GPU, pass
452452
# ``torch.device('cpu')`` to the ``map_location`` argument in the
@@ -469,7 +469,7 @@
469469
#
470470
# device = torch.device("cuda")
471471
# model = TheModelClass(*args, **kwargs)
472-
# model.load_state_dict(torch.load(PATH))
472+
# model.load_state_dict(torch.load(PATH, weights_only=True))
473473
# model.to(device)
474474
# # Make sure to call input = input.to(device) on any input tensors that you feed to the model
475475
#
@@ -497,7 +497,7 @@
497497
#
498498
# device = torch.device("cuda")
499499
# model = TheModelClass(*args, **kwargs)
500-
# model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose whatever GPU device number you want
500+
# model.load_state_dict(torch.load(PATH, weights_only=True, map_location="cuda:0")) # Choose whatever GPU device number you want
501501
# model.to(device)
502502
# # Make sure to call input = input.to(device) on any input tensors that you feed to the model
503503
#

beginner_source/transfer_learning_tutorial.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
209209
print(f'Best val Acc: {best_acc:4f}')
210210

211211
# load best model weights
212-
model.load_state_dict(torch.load(best_model_params_path))
212+
model.load_state_dict(torch.load(best_model_params_path, weights_only=True))
213213
return model
214214

215215

intermediate_source/autograd_saved_tensors_hooks_tutorial.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -397,7 +397,7 @@ def pack_hook(tensor):
397397
return name
398398

399399
def unpack_hook(name):
400-
return torch.load(name)
400+
return torch.load(name, weights_only=True)
401401

402402

403403
######################################################################
@@ -420,7 +420,7 @@ def pack_hook(tensor):
420420
return name
421421

422422
def unpack_hook(name):
423-
tensor = torch.load(name)
423+
tensor = torch.load(name, weights_only=True)
424424
os.remove(name)
425425
return tensor
426426

@@ -462,7 +462,7 @@ def pack_hook(tensor):
462462
return temp_file
463463

464464
def unpack_hook(temp_file):
465-
return torch.load(temp_file.name)
465+
return torch.load(temp_file.name, weights_only=True)
466466

467467

468468
######################################################################

intermediate_source/ddp_tutorial.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,7 @@ and elasticity support, please refer to `TorchElastic <https://pytorch.org/elast
214214
# configure map_location properly
215215
map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
216216
ddp_model.load_state_dict(
217-
torch.load(CHECKPOINT_PATH, map_location=map_location))
217+
torch.load(CHECKPOINT_PATH, map_location=map_location, weights_only=True))
218218
219219
loss_fn = nn.MSELoss()
220220
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

0 commit comments

Comments
 (0)