pytorch使用matplotlib绘制神经网络模型的loss（损失值）曲线

基于这篇博客中遇到的显示x和y维度不一致问题（Python可视化）使用matplotlib.pyplot绘制神经网络模型的acc(准确率)和loss(损失值)曲线图_lstm模型loss与accuracy画图-CSDN博客

实验发现是因为保存的loss为跑完每一个batch以后的total loss，将total loss平均得到loss，注意要在每一轮结束后把平均的loss添加到列表里而不是在每一个batch结束后将平均的loss填入list，简单来说就是把Loss_list.append(f'{total_train_loss:.2f}')添加到哪个循环内的问题，为了保持维度一致需要添加在对epoch的外循环里。

完整代码如下：

import math
import warnings
import numpy as np
import torch.optim
from sklearn.metrics import confusion_matrix
from sklearn.preprocessing import LabelEncoder
import pandas as pd
from torch import nn
from torch.utils.data import DataLoader
from matplotlib import pyplot as plt

warnings.simplefilter(action='ignore', category=RuntimeWarning)
torch.manual_seed(2022)
#下面的mynet是随便写的一个模型，把这部分改成自己的模型就可以
class mynet(nn.Module):
    def __init__(self):
        super(mynet, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv1d( 1, 16, 2),
            nn.ReLU(),
            nn.MaxPool1d(4), 
            nn.Flatten(), 
        )
        self.model2 = nn.Sequential(
            nn.Linear(in_features=32, out_features=3, bias=True),#一个3分类的数据集，输出特征为3
            nn.Softmax(dim=1)#softmax得到在每个类别上的概率
        )

    def forward(self, input):
        input = input.reshape(-1, 1, 6)  
        x = self.model1(input)
        x = self.model2(x)
        return x

data = pd.read_csv('train.csv')#自己的数据集地址
dataset = data.values
X_train = dataset[:, 1:7].astype(float)#1到6列为特征
Y_train = dataset[:, 0:1]#0列为标签
encoder = LabelEncoder()
Y_train = encoder.fit_transform(Y_train.ravel())
X_train, Y_train = torch.FloatTensor(X_train), torch.LongTensor(Y_train)
train_dataset = torch.utils.data.TensorDataset(X_train, Y_train)#这里就只写train数据集

loss_function = nn.CrossEntropyLoss()#损失函数为交叉熵损失
epochs = 100
learning_rate = 0.001
optim = torch.optim.Adam(mynet.parameters(), lr=learning_rate)

Loss_list = [] #注意创建的loss_list要放在每一轮训练的前面
for i in range(epochs):#外循环是epoch的循环
    print("--------第{}轮训练开始---------".format(i+1))
    total_train_loss = 0
    train_loader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
    mynet.train()
    for data in train_loader:#内循环是对batch的循环
        X_data, Y_data = data[0], data[1]
        output = mynet(X_data)
        train_loss = loss_function(output, Y_data)
        optim.zero_grad()
        train_loss.backward()
        optim.step()
        total_train_loss = total_train_loss + train_loss#注意这里得到的total_loss是内循环中每一个batch的loss加和
    length = len(train_loader)
    train_loss = total_train_loss / length#最终填入的loss要除以train的length，因为前面得到的是所有batch的总loss，得到一个loss的平均
    Loss_list.append(f'{train_loss:.2f}')#循环将每一轮的train_loss填入list
    print("train_loss为{:.4f}".format(train_loss))
 

# 对存入列表中的数据进行强制类型转换（转换的类型大家自选，这里是float）
for i in range(0, len(Loss_list)):
    Loss_list[i] = float(Loss_list[i])
# x轴数据的取值范围（训练多少次就填多少）
x = range(0, 100)
y = Loss_list
# loss图像
plt.plot(x, y, '.-')
plt.xlabel('model loss')
plt.ylabel('loss')
my_yTicks2 = np.arange(0.2, 2, 0.2)
plt.yticks(my_yTicks2)
plt.savefig("loss.jpg")
plt.show()