创建:
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张量