pyTorch 保存提取模型
import torch
# 保存
torch.save(the_model, PATH)
# 恢复
the_model = torch.load(PATH)
# 该方法保存的数据绑定着特定的 classes 和所用的确切目录结构.
‘ # 因此，再加载后经过许多重构后，可能会被打乱.
#### 官网推荐的方法
保存
torch.save(the_model.state_dict(), PATH)
恢复
the_model = TheModelClass(*args, **kwargs) ### 输入神经网络的层数
the_model.load_state_dict(torch.load(PATH))
使用这种方法，我们需要自己导入模型的结构信息。调用：
import numpy as np
import scipy.io as scio
import torch
import torch.nn.functional as Fun
import torch.nn as nn
import torchvision
import mat4py
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
class TestDataset(Dataset):
def __init__(self):
dataFile='./test_data_5_302_0318.mat'
test = mat4py.loadmat(dataFile)['test_data']
self.train_dataset=test### 76943*343
self.data_tf = transforms.ToTensor()
def __getitem__(self, i):
return torch.from_numpy(np.array(self.train_dataset[i][0:302])), torch.from_numpy(np.array([self.train_dataset[i][302:]]))
def __len__(self):
return len(self.train_dataset)
test_dataset = TestDataset( )
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=128,
shuffle=True)
class BiRNN(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, num_classes):
super(BiRNN, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
self.fc = nn.Linear(hidden_size*2, num_classes) # 2 for bidirection
def forward(self, x):
# Set initial states
h0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size) # 2 for bidirection
c0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size)
# Forward propagate LSTM
out, _ = self.lstm(x, (h0, c0)) # out: tensor of shape (batch_size, seq_length, hidden_size*2)
# Decode the hidden state of the last time step
out = self.fc(out[:, -1, :])
return out
##torch.save(model.state_dict(), 'model_bilstm_v2.ckpt') save model
model = BiRNN(302,1024,4,5)
model.load_state_dict(torch.load('model_bilstm_v2.ckpt'))
#print("model",model,type(model))
y_true1=[]
y_pred1=[]
with torch.no_grad():
correct = 0
total = 0
for data, labels in test_loader:
labels=labels.view(labels.shape[0],labels.shape[2])
data=data.view(data.shape[0],1,data.shape[1])
outputs = model(data.float())
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
labels_= torch.max(labels, 1)[1]
correct += (predicted == labels_).sum().item()
predicted=predicted.cpu().numpy().tolist()
labels_=labels_.cpu().numpy().tolist()
y_true1.extend(predicted)
y_pred1.extend(labels_)
print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))
#torch.save(model.state_dict(), 'model_bilstm_v2.ckpt')
y_true1=np.array(y_true1)
y_pred1=np.array(y_pred1)
a=confusion_matrix(y_pred1,y_true1 , labels=[0,1,2,3,4])
print("a:",a)
print("**************")
total=sum(sum(a))#### important value of sum
print(total,a.shape)
##all acc:
total_OA=0
for i in range(a.shape[0]):
total_OA+=a[i][i]
print("all_acc",total_OA/total)
## category acc
ca_total=[0,0,0,0,0]
m=0
for k in range(5):
for i in range(a.shape[0]):
for j in range(a.shape[0]):
if i!=k and j!=k:
# print(a[i][j])
ca_total[k]+=a[i][j]
ca_total[k]+=a[k][k]
ca_total[k]=ca_total[k]/total
print(" five classifity acc:",ca_total)
## sensitivity
Se=[0,0,0,0,0]
print(sum(a[0,:]))
for i in range(5):
Se[i]=a[i][i]/sum(a[i,:])
print("five calssify sensitivity:",Se)
### Specificity
Sp=[0,0,0,0,0]
Sp_1=[0,0,0,0,0]
for k in range(5):
for i in range(5):
for j in range(5):
if i!=k and j!=k:
Sp_1[k]+=a[i][j]
if i!=k and j==k:
Sp[k]+=a[i][j]
Sp[k]=Sp_1[k]/(Sp_1[k]+Sp[k])
print("Sp:",Sp)
### Positive Predicti
Pp=[0,0,0,0,0]
for i in range(5):
Pp[i]=a[i][i]/sum(a[:,i])
print("Pp:",Pp)
print("model:",model,num_epochs,batchsize,learn_rate)