PyTorch神经网络模型从CPU或GPU中加载和保存

保存模型，分别保存CPU模型和GPU模型:

代码展示:

# 分别为来自CPU和GPU的模型保存参数，
# 并再次将模型的参数加载到位于CPU或者GPU的模型上

import torch
import torch.nn as nn
torch.manual_seed(seed=20200910)
class Model(torch.nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1=torch.nn.Sequential(  # 输入torch.Size([64, 1, 28, 28])
                torch.nn.Conv2d(1,64,kernel_size=3,stride=1,padding=1),
                torch.nn.ReLU(),  # 输出torch.Size([64, 64, 28, 28])
                torch.nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),  # 输出torch.Size([64, 128, 28, 28])
                torch.nn.ReLU(),
                torch.nn.MaxPool2d(stride=2,kernel_size=2)  # 输出torch.Size([64, 128, 14, 14])
        )

        self.dense=torch.nn.Sequential(  # 输入torch.Size([64, 14*14*128])
                    torch.nn.Linear(14*14*128,1024),  # 输出torch.Size([64, 1024])
                    torch.nn.ReLU(),
                    torch.nn.Dropout(p=0.5),
                    torch.nn.Linear(1024,10)  # 输出torch.Size([64, 10])        
        )
        self.layer4cxq1 = torch.nn.Conv2d(2,33,4,4)
        self.layer4cxq2 = torch.nn.ReLU()
        self.layer4cxq3 = torch.nn.MaxPool2d(stride=2,kernel_size=2)
        self.layer4cxq4 = torch.nn.Linear(14*14*128,1024)
        self.layer4cxq5 = torch.nn.Dropout(p=0.8)
        self.attribute4cxq = nn.Parameter(torch.tensor(20200910.0))
        self.attribute4lzq = nn.Parameter(torch.tensor([2.0,3.0,4.0,5.0]))
        self.attribute4hh = nn.Parameter(torch.randn(3,4,5,6))
        self.attribute4wyf = nn.Parameter(torch.randn(7,8,9,10))

    def forward(self,x):  # torch.Size([64, 1, 28, 28])
        x = self.conv1(x)  # 输出torch.Size([64, 128, 14, 14])
        x = x.view(-1,14*14*128)  # torch.Size([64, 14*14*128])
        x = self.dense(x)  # 输出torch.Size([64, 10])
        return x

print('cuda(GPU)是否可用:',torch.cuda.is_available())  # cuda(GPU)是否可用: True
print('torch的版本:',torch.__version__)  # torch的版本: 1.2.0+cu92
print('\n\n\n')

model = Model().cpu() #.cuda()
print(len(list(model.parameters())))  # 16
print('模型所在位置:',next(model.parameters()).device)  # cpu
print(len(list(model.parameters())))  #  16

print('保存的参数位于CPU上的模型'.center(100,'-'))

state_dict = model.state_dict()
model_params_path = 'CPU参数/模型参数.pth'
torch.save(state_dict, model_params_path)
print((model_params_path+'保存成功!!!').center(100,'-'))



print('\n\n\n')


model = model.cuda()
print('模型所在位置:',next(model.parameters()).device)  # cuda:0
print('保存的参数位于GPU上的模型'.center(100,'-'))
state_dict = model.state_dict()
model_params_path = 'GPU参数/模型参数.pth'
torch.save(state_dict, model_params_path)
print((model_params_path+'保存成功!!!').center(100,'-'))

这里为了删除某一层神经网络，更严谨的话应该使用del关键字，而不是简单地将对象的属性设置为None;完整修改代码如下:

# 分别为来自CPU和GPU的模型保存参数，
# 并再次将模型的参数加载到位于CPU或者GPU的模型上

import torch
import torch.nn as nn
torch.manual_seed(seed=20200910)
class Model(torch.nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1=torch.nn.Sequential(  # 输入torch.Size([64, 1, 28, 28])
                torch.nn.Conv2d(1,64,kernel_size=3,stride=1,padding=1),
                torch.nn.ReLU(),  # 输出torch.Size([64, 64, 28, 28])
                torch.nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),  # 输出torch.Size([64, 128, 28, 28])
                torch.nn.ReLU(),
                torch.nn.MaxPool2d(stride=2,kernel_size=2)  # 输出torch.Size([64, 128, 14, 14])
        )

        self.dense=torch.nn.Sequential(  # 输入torch.Size([64, 14*14*128])
                    torch.nn.Linear(14*14*128,1024),  # 输出torch.Size([64, 1024])
                    torch.nn.ReLU(),
                    torch.nn.Dropout(p=0.5),
                    torch.nn.Linear(1024,10)  # 输出torch.Size([64, 10])        
        )
        self.layer4cxq1 = torch.nn.Conv2d(2,33,4,4)
        self.layer4cxq2 = torch.nn.ReLU()

        self.layer4cxq3 = torch.nn.MaxPool2d(stride=2,kernel_size=2)
        self.layer4cxq4 = torch.nn.Linear(14*14*128,1024)
        self.layer4cxq5 = torch.nn.Dropout(p=0.8)

        self.attribute4cxq = nn.Parameter(torch.tensor(20200910.0))
        self.attribute4lzq = nn.Parameter(torch.tensor([2.0,3.0,4.0,5.0]))

        self.attribute4hh = nn.Parameter(torch.randn(3,4,5,6))
        self.attribute4wyf = nn.Parameter(torch.randn(7,8,9,10))

    def forward(self,x):  # torch.Size([64, 1, 28, 28])
        x = self.conv1(x)  # 输出torch.Size([64, 128, 14, 14])
        x = x.view(-1,14*14*128)  # torch.Size([64, 14*14*128])
        x = self.dense(x)  # 输出torch.Size([64, 10])
        return x

    def delete_fore_half(self):
        del self.conv1
        del self.layer4cxq1
        del self.layer4cxq2
        del self.attribute4cxq
        del self.attribute4lzq
        return self

    def delete_back_half(self):
        del self.dense
        del self.layer4cxq3
        del self.layer4cxq4
        del self.layer4cxq5
        del self.attribute4hh
        del self.attribute4wyf
        return self

    def delete_to_A(self):
        del self.layer4cxq4
        del self.layer4cxq5
        del self.attribute4wyf
        return self

    def delete_to_B(self):
        del self.dense
        del self.layer4cxq3
        del self.attribute4hh
        return self
if __name__ == '__main__':
    print('cuda(GPU)是否可用:',torch.cuda.is_available())  # cuda(GPU)是否可用: True
    print('torch的版本:',torch.__version__)  # torch的版本: 1.2.0+cu92
    print('\n')

    model = Model()
    print('完整模型:')
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model_params_path = '完整模型.pth'
    torch.save(model.state_dict(), model_params_path)



    print()
    print('模型删除前半部分'.center(100,'-'))
    print('删除前:')
    model = Model()
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model = model.delete_fore_half()
    print('删除后:')
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model_params_path = '后半个模型.pth'
    torch.save(model.state_dict(), model_params_path)






    print()
    print('模型删除后半部分'.center(100,'-'))
    print('删除前:')
    model = Model()
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model = model.delete_back_half()
    print('删除后:')
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model_params_path = '前半个模型.pth'
    torch.save(model.state_dict(), model_params_path)




    print()
    print('模型删除到模型A'.center(100,'-'))
    print('删除前:')
    model = Model()
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model = model.delete_to_A()
    print('删除后:')
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model_params_path = '模型A部分.pth'
    torch.save(model.state_dict(), model_params_path)




    print()
    print('模型删除到模型B'.center(100,'-'))
    print('删除前:')
    model = Model()
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model = model.delete_to_B()
    print('删除后:')
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    print('模型所在位置:',next(model.parameters()).device)
    print('模型中参数数量:',len(list(model.parameters())))
    print('模型中状态字典长度:',len(model.state_dict()))
    model_params_path = '模型B部分.pth'
    torch.save(model.state_dict(), model_params_path)

执行后效果截图:

控制台输出结果展示:

版权所有 (C) Microsoft Corporation。保留所有权利。

尝试新的跨平台 PowerShell https://aka.ms/pscore6

加载个人及系统配置文件用了 1033 毫秒。
(base) PS C:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)>  & 'D:\Anaconda3\python.exe' 'c:\Users\chenxuqi\.vscode\extensions\ms-python.python-2021.1.502429796\pythonFiles\lib\python\debugpy\launcher' '54056' '--' 'c:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)\test1_保存模型.py'
cuda(GPU)是否可用: True
torch的版本: 1.2.0+cu92




16
模型所在位置: cpu
16
-------------------------------------------保存的参数位于CPU上的模型-------------------------------------------
---------------------------------------CPU参数/模型参数.pth保存成功!!!----------------------------------------




模型所在位置: cuda:0
-------------------------------------------保存的参数位于GPU上的模型-------------------------------------------
---------------------------------------GPU参数/模型参数.pth保存成功!!!----------------------------------------
(base) PS C:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)> conda activate base
(base) PS C:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)>

CPU和GPU之间的数据相互加载:
代码展示:

# 分别为来自CPU和GPU的模型保存参数，
# 并再次将模型的参数加载到位于CPU或者GPU的模型上

import torch
import torch.nn as nn
torch.manual_seed(seed=20200910)
class Model(torch.nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1=torch.nn.Sequential(  # 输入torch.Size([64, 1, 28, 28])
                torch.nn.Conv2d(1,64,kernel_size=3,stride=1,padding=1),
                torch.nn.ReLU(),  # 输出torch.Size([64, 64, 28, 28])
                torch.nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),  # 输出torch.Size([64, 128, 28, 28])
                torch.nn.ReLU(),
                torch.nn.MaxPool2d(stride=2,kernel_size=2)  # 输出torch.Size([64, 128, 14, 14])
        )

        self.dense=torch.nn.Sequential(  # 输入torch.Size([64, 14*14*128])
                    torch.nn.Linear(14*14*128,1024),  # 输出torch.Size([64, 1024])
                    torch.nn.ReLU(),
                    torch.nn.Dropout(p=0.5),
                    torch.nn.Linear(1024,10)  # 输出torch.Size([64, 10])        
        )
        self.layer4cxq1 = torch.nn.Conv2d(2,33,4,4)
        self.layer4cxq2 = torch.nn.ReLU()
        self.layer4cxq3 = torch.nn.MaxPool2d(stride=2,kernel_size=2)
        self.layer4cxq4 = torch.nn.Linear(14*14*128,1024)
        self.layer4cxq5 = torch.nn.Dropout(p=0.8)
        self.attribute4cxq = nn.Parameter(torch.tensor(20200910.0))
        self.attribute4lzq = nn.Parameter(torch.tensor([2.0,3.0,4.0,5.0]))
        self.attribute4hh = nn.Parameter(torch.randn(3,4,5,6))
        self.attribute4wyf = nn.Parameter(torch.randn(7,8,9,10))

    def forward(self,x):  # torch.Size([64, 1, 28, 28])
        x = self.conv1(x)  # 输出torch.Size([64, 128, 14, 14])
        x = x.view(-1,14*14*128)  # torch.Size([64, 14*14*128])
        x = self.dense(x)  # 输出torch.Size([64, 10])
        return x

print('cuda(GPU)是否可用:',torch.cuda.is_available())  # cuda(GPU)是否可用: True
print('torch的版本:',torch.__version__)  # torch的版本: 1.2.0+cu92
print('\n\n\n')

model = Model()
print('模型所在位置:',next(model.parameters()).device)

model = Model().cpu()
print('模型所在位置:',next(model.parameters()).device)
print()

# -----------------------CPU->CPU-------------------------------- #
print('从CPU加载到CPU'.center(100,'-'))
model_params_path = 'CPU参数/模型参数.pth'
model = Model()
print('模型所在位置:',next(model.parameters()).device)
model.load_state_dict(torch.load(model_params_path))
print('模型所在位置:',next(model.parameters()).device)
print()



# -----------------------GPU->GPU-------------------------------- #
print('从GPU加载到GPU'.center(100,'-'))
model_params_path = 'GPU参数/模型参数.pth'
model = Model().cuda()
print('模型所在位置:',next(model.parameters()).device)
model.load_state_dict(torch.load(model_params_path))
print('模型所在位置:',next(model.parameters()).device)
print()






# -----------------------CPU->GPU-------------------------------- #
print('从CPU加载到GPU'.center(100,'-'))
model_params_path = 'CPU参数/模型参数.pth'
model = Model().cuda()
print('模型所在位置:',next(model.parameters()).device)
model.load_state_dict(torch.load(model_params_path))
print('模型所在位置:',next(model.parameters()).device)
print()





# -----------------------GPU->CPU-------------------------------- #
print('从GPU加载到CPU'.center(100,'-'))
model_params_path = 'GPU参数/模型参数.pth'
model = Model().cpu()
print('模型所在位置:',next(model.parameters()).device)
model.load_state_dict(torch.load(model_params_path))
print('模型所在位置:',next(model.parameters()).device)
print()

控制台输出结果展示:

Windows PowerShell
版权所有 (C) Microsoft Corporation。保留所有权利。

尝试新的跨平台 PowerShell https://aka.ms/pscore6

加载个人及系统配置文件用了 876 毫秒。
(base) PS C:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)> conda activate pytorch_1.7.1_cu102
(pytorch_1.7.1_cu102) PS C:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)>  & 'D:\Anaconda3\envs\pytorch_1.7.1_cu102\python.exe' 'c:\Users\chenxuqi\.vscode\extensions\ms-python.python-2021.1.502429796\pythonFiles\lib\python\debugpy\launcher' '54551' '--' 'c:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)\test2_加载模型.py'
cuda(GPU)是否可用: True
torch的版本: 1.7.1




模型所在位置: cpu
模型所在位置: cpu

---------------------------------------------从CPU加载到CPU---------------------------------------------
模型所在位置: cpu
模型所在位置: cpu

---------------------------------------------从GPU加载到GPU---------------------------------------------
模型所在位置: cuda:0
模型所在位置: cuda:0

---------------------------------------------从CPU加载到GPU---------------------------------------------
模型所在位置: cuda:0
模型所在位置: cuda:0

---------------------------------------------从GPU加载到CPU---------------------------------------------
模型所在位置: cpu
模型所在位置: cpu

(pytorch_1.7.1_cu102) PS C:\Users\chenxuqi\Desktop\News4cxq\测试模型保存与加载(GPU和CPU)>

程序执行效果:

本文同步分享在 博客“敲代码的小风”（CSDN）。
如有侵权，请联系 [email protected] 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。