torch模型创建、保存与加载

创建：
class 类名()：
def int(self):
super(类名，self).init()

def forward(self，input):#执行

#保存
path=’’.pt’’
torch.save()
#加载
model='模型名称‘’
torch.model.load_state_dict()

import torch
import torch.nn as nn
#卷积-drop-展平-减少特征数linear
#定义模型
class NeutralModel(nn.Module):
  def __init__(self):
    super(NeutralModel,self).__init__()
    self.conv1=nn.Conv2d(1,32,3,1)
    # self.conv2=nn.Conv2d(16,8,3,1)
    self.drop=nn.Dropout(0.3)
    self.fun=nn.Flatten()
    self.line=nn.Linear(21632,10)
  def forward(self,x):
    x=self.conv1(x)
    x=nn.functional.relu(x)
    # x=self.conv2(x)
    x=self.drop(x)
    x=self.fun(x)
    x=self.line(x)
    return x

input_data=torch.rand(1,1,28,28)#相当于一张随机的 28x28 图像
print(input_data)
model=NeutralModel()
result=model(input_data)
print(result)
print(result.shape)


optimizer=torch.optim.SGD(model.parameters(),lr=0.1,momentum=0.9)
#访问模型和优化器
print('模型的sate_dict:')
print(model.state_dict().keys())
for param_tensor in model.state_dict():
  print(param_tensor,'\t',model.state_dict()[param_tensor].size())
print('优化器的state_dict：')
for var in optimizer.state_dict():
  print(var,'\t',optimizer.state_dict()[var])



#保存和加载模型
path='model.pt'
torch.save(model.state_dict(),path)
#加载
device=torch.device('cpu')#在cpu上保存
# device=torch.device('cuda')#在GPU上保存
model=NeutralModel()
model.load_state_dict(torch.load(path,map_location=device))


#在GPU上加载模型
path='model.pt'
torch.save(model.state_dict(),path)
device=torch.device('cuda')
model=NeutralModel()
model.load_state_dict(torch.load(path,map_location='cuda:0'))#分配给指定的GPU设备
model.to(device)#转为cuda张量