自己做网站去哪买服务器,改网站描述,当今做那些网站能致富,灵犀科技 高端网站建设PyTorch 模型保存读取
模型存储模型单卡存储多卡存储模型单卡读取多卡读取
1、模型存储
PyTorch存储模型主要采用pkl#xff0c;pt#xff0c;pth三种格式,就使用层面来说没有区别PyTorch模型主要包含两个部分#xff1a;模型结构和权重。其中模型是继承n…PyTorch 模型保存读取
模型存储模型单卡存储多卡存储模型单卡读取多卡读取
1、模型存储
PyTorch存储模型主要采用pklptpth三种格式,就使用层面来说没有区别PyTorch模型主要包含两个部分模型结构和权重。其中模型是继承nn.Module的类权重的数据结构是一个字典key是层名value是权重向量存储也由此分为两种形式存储整个模型包括结构和权重和只存储模型权重(推荐。
import torch
from torchvision import models
model models.resnet50(pretrainedTrue)
save_dir ./resnet50.pth# 保存整个 模型结构权重
torch.save(model, save_dir)
# 保存 模型权重
torch.save(model.state_dict, save_dir)# pt, pth和pkl三种数据格式均支持模型权重和整个模型的存储2、模型单卡存储多卡存储
PyTorch中将模型和数据放到GPU上有两种方式——.cuda()和.to(device)注如果要使用多卡训练的话需要对模型使用torch.nn.DataParallel
2.1、nn.DataParrallel
CLASS torch.nn.DataParallel(module, device_idsNone, output_deviceNone, dim0) [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-puyISgkD-1692613764220)(attachment:image.png)]
module即表示你定义的模型device_ids表示你训练的deviceoutput_device这个参数表示输出结果的device而这最后一个参数output_device一般情况下是省略不写的那么默认就是在device_ids[0]
注因此一般情况下第一张显卡的内存使用占比会更多
import os
import torch
from torchvision import models#单卡
os.environ[CUDA_VISIBLE_DEVICES] 0 # 如果是多卡改成类似0,1,2
model model.cuda() # 单卡
#print(model)---------------------------------------------------------------------------RuntimeError Traceback (most recent call last)~\AppData\Local\Temp/ipykernel_7460/77570021.py in module1 import os2 os.environ[CUDA_VISIBLE_DEVICES] 0 # 如果是多卡改成类似0,1,2
---- 3 model model.cuda() # 单卡D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in cuda(self, device)903 Module: self904
-- 905 return self._apply(lambda t: t.cuda(device))906 907 def ipu(self: T, device: Optional[Union[int, device]] None) - T:D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)795 def _apply(self, fn):796 for module in self.children():
-- 797 module._apply(fn)798 799 def compute_should_use_set_data(tensor, tensor_applied):D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)818 # with torch.no_grad():819 with torch.no_grad():
-- 820 param_applied fn(param)821 should_use_set_data compute_should_use_set_data(param, param_applied)822 if should_use_set_data:D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in lambda(t)903 Module: self904
-- 905 return self._apply(lambda t: t.cuda(device))906 907 def ipu(self: T, device: Optional[Union[int, device]] None) - T:D:\Users\xulele\Anaconda3\lib\site-packages\torch\cuda\__init__.py in _lazy_init()245 if CUDA_MODULE_LOADING not in os.environ:246 os.environ[CUDA_MODULE_LOADING] LAZY
-- 247 torch._C._cuda_init()248 # Some of the queued calls may reentrantly call _lazy_init();249 # we need to just return without initializing in that case.RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-0G4NTv1z-1692613764220)(attachment:ed8eb711294e4c6e3e43690ddb2bf66.png)]
#多卡
os.environ[CUDA_VISIBLE_DEVICES] 0,1
model torch.nn.DataParallel(model).cuda() # 多卡
#print(model)[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-eHt1Dn8t-1692613764221)(attachment:image.png)]
2.3、单卡保存单卡加载
os.environ[CUDA_VISIBLE_DEVICES] 0 #这里替换成希望使用的GPU编号
model models.resnet50(pretrainedTrue)
model.cuda()save_dir resnet50.pt #保存路径# 保存读取整个模型
torch.save(model, save_dir)
loaded_model torch.load(save_dir)
loaded_model.cuda()# 保存读取模型权重
torch.save(model.state_dict(), save_dir)
# 先加载模型结构
loaded_model models.resnet50()
# 在加载模型权重
loaded_model.load_state_dict(torch.load(save_dir))
loaded_model.cuda()---------------------------------------------------------------------------RuntimeError Traceback (most recent call last)~\AppData\Local\Temp/ipykernel_7460/585340704.py in module5 os.environ[CUDA_VISIBLE_DEVICES] 0 #这里替换成希望使用的GPU编号6 model models.resnet50(pretrainedTrue)
---- 7 model.cuda()8 9 save_dir resnet50.pt #保存路径D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in cuda(self, device)903 Module: self904
-- 905 return self._apply(lambda t: t.cuda(device))906 907 def ipu(self: T, device: Optional[Union[int, device]] None) - T:D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)795 def _apply(self, fn):796 for module in self.children():
-- 797 module._apply(fn)798 799 def compute_should_use_set_data(tensor, tensor_applied):D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)818 # with torch.no_grad():819 with torch.no_grad():
-- 820 param_applied fn(param)821 should_use_set_data compute_should_use_set_data(param, param_applied)822 if should_use_set_data:D:\Users\xulele\Anaconda3\lib\site-packages\torch\nn\modules\module.py in lambda(t)903 Module: self904
-- 905 return self._apply(lambda t: t.cuda(device))906 907 def ipu(self: T, device: Optional[Union[int, device]] None) - T:D:\Users\xulele\Anaconda3\lib\site-packages\torch\cuda\__init__.py in _lazy_init()245 if CUDA_MODULE_LOADING not in os.environ:246 os.environ[CUDA_MODULE_LOADING] LAZY
-- 247 torch._C._cuda_init()248 # Some of the queued calls may reentrantly call _lazy_init();249 # we need to just return without initializing in that case.RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx2.4、单卡保存多卡加载 os.environ[CUDA_VISIBLE_DEVICES] 0 #这里替换成希望使用的GPU编号
model models.resnet50(pretrainedTrue)
model.cuda()# 保存读取整个模型
torch.save(model, save_dir)os.environ[CUDA_VISIBLE_DEVICES] 1,2 #这里替换成希望使用的GPU编号
loaded_model torch.load(save_dir)
loaded_model nn.DataParallel(loaded_model).cuda()# 保存读取模型权重
torch.save(model.state_dict(), save_dir)os.environ[CUDA_VISIBLE_DEVICES] 1,2 #这里替换成希望使用的GPU编号
loaded_model models.resnet50() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir))
loaded_model nn.DataParallel(loaded_model).cuda()2.5、多卡保存单卡加载
核心问题如何去掉权重字典键名中的module以保证模型的统一性
对于加载整个模型直接提取模型的module属性即可对于加载模型权重保存模型时保存模型的module属性对应的权重
os.environ[CUDA_VISIBLE_DEVICES] 1,2 #这里替换成希望使用的GPU编号model models.resnet50(pretrainedTrue)
model nn.DataParallel(model).cuda()# 保存读取整个模型
torch.save(model, save_dir)os.environ[CUDA_VISIBLE_DEVICES] 0 #这里替换成希望使用的GPU编号
loaded_model torch.load(save_dir).moduleos.environ[CUDA_VISIBLE_DEVICES] 0,1,2 #这里替换成希望使用的GPU编号model models.resnet50(pretrainedTrue)
model nn.DataParallel(model).cuda()# 保存权重
torch.save(model.module.state_dict(), save_dir)#加载模型权重
os.environ[CUDA_VISIBLE_DEVICES] 0 #这里替换成希望使用的GPU编号
loaded_model models.resnet50() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir))
loaded_model.cuda()
2.6、多卡保存多卡加载
保存整个模型时会同时保存所使用的GPU id等信息读取时若这些信息和当前使用的GPU信息不符则可能会报错或者程序不按预定状态运行。可能出现以下2个问题
1、读取整个模型再使用nn.DataParallel进行分布式训练设置这种情况很可能会造成保存的整个模型中GPU id和读取环境下设置的GPU id不符训练时数据所在device和模型所在device不一致而报错2、读取整个模型而不使用nn.DataParallel进行分布式训练设置发现程序会自动使用设备的前n个GPU进行训练n是保存的模型使用的GPU个数。此时如果指定的GPU个数少于n则会报错
建议方案
只模型权重之后再使用nn.DataParallel进行分布式训练设置则没有问题因此多卡模式下建议使用权重的方式存储和读取模型
os.environ[CUDA_VISIBLE_DEVICES] 0,1,2 #这里替换成希望使用的GPU编号model models.resnet50(pretrainedTrue)
model nn.DataParallel(model).cuda()# 保存读取模型权重强烈建议
torch.save(model.state_dict(), save_dir)
#加载模型 权重
loaded_model models.resnet50() #注意这里需要对模型结构有定义
loaded_model.load_state_dict(torch.load(save_dir)))
loaded_model nn.DataParallel(loaded_model).cuda()建议
不管是单卡保存还是多卡保存建议以保存模型权重为主不管是单卡还是多卡,先load模型权重再指定是多卡加载nn.DataParallel或单卡(cuda)
# 使用案例截取片段代码My_model.eval()
test_total_loss 0
test_total_correct 0
test_total_num 0past_test_loss 0 #上一轮的loss
save_model_step 10 # 每10步保存一次modelfor iter,(images,labels) in enumerate(test_loader):images images.to(device)labels labels.to(device)outputs My_model(images)loss criterion(outputs,labels)test_total_correct (outputs.argmax(1) labels).sum().item()test_total_loss loss.item()test_total_num labels.shape[0]test_loss test_total_loss / test_total_numprint(Epoch [{}/{}], train_loss:{:.4f}, train_acc:{:.4f}%, test_loss:{:.4f}, test_acc:{:.4f}%.format(i1, epoch, train_total_loss / train_total_num, train_total_correct / train_total_num * 100, test_total_loss / test_total_num, test_total_correct / test_total_num * 100))# model saveif test_losspast_test_loss:#保存模型权重torch.save(model.state_dict(), save_dir)#保存 模型权重模型结构#torch.save(model, save_dir)if iter % save_model_step 0:#保存模型权重torch.save(model.state_dict(), save_dir)#保存 模型权重模型结构#torch.save(model, save_dir)past_test_loss test_loss
单卡保存单卡读取 案例
Google Colab:https://colab.research.google.com/drive/1hEOeqXYm4BfulY6d30QCI4HrFmCmmTQu?uspsharing
