pyTorch 保存提取模型

水深无声 2022-04-05 14:21 262阅读 0赞

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)