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)