深度学习代码实践_train.py文件内容（识别数字0-9）

import cv2
from MLP import MLP
from Cnn import save_model
import torch
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.model_selection import train_test_split
import torch.utils.data as Data
import numpy as np
import matplotlib.pyplot as plt
from data_process import get_features
from sklearn.preprocessing import StandardScaler

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

train_losses = []
train_accuracy = []
def train_model(net, data, label, lr, batch_size, epoch):
    print(net)
    data = torch.Tensor(data)
    data = data.unsqueeze(1)
    label = torch.Tensor(label).long()

    #resnet建议使用cuda
    data = data.to(device)
    label = label.to(device)
    # 训练集和测试集7：3
    train_data, test_data, train_label, test_label = train_test_split(data, label, test_size=0.3, random_state=0)

    # 学习率
    LR = lr
    # 每次投入训练数据大小
    BATCH_SIZE = batch_size
    # 训练模型次数
    EPOCH = epoch

    optimizer = torch.optim.Adam(net.parameters(), lr=LR)

    train_dataset = Data.TensorDataset(train_data, train_label)
    train_loader = Data.DataLoader(
        dataset=train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
    )

    test_dataset = Data.TensorDataset(test_data, test_label)
    test_loader = Data.DataLoader(
        dataset=test_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
    )
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, LR, epochs=EPOCH, steps_per_epoch=len(train_loader))


    for epoch in range(EPOCH):
        running_loss = 0
        for step, (batch_data, batch_label) in enumerate(train_loader):
            print('Epoch:', epoch + 1, '/', EPOCH, 'Step:', step)
            prediction = net(batch_data)
            loss = F.cross_entropy(prediction, batch_label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()
            running_loss += loss.item()

            _, pred = torch.max(prediction, 1)

            accuracy = torch.sum(pred == batch_label).item() / len(pred)
            print('Epoch', epoch + 1, '| train loss:%.4f' % loss, '| accuracy:%.4f' % accuracy)
        train_losses.append(running_loss / len(train_loader))
        train_accuracy.append(accuracy)

    return net


def test_model(net, data, label):
    net.eval()
    data = torch.Tensor(data)
    data = data.unsqueeze(1)
    label = torch.Tensor(label).long()
    data = data.to(device)
    label = label.to(device)
    # 训练集和测试集7：3
    train_data, test_data, train_label, test_label = train_test_split(data, label, test_size=0.3, random_state=0)

    test_dataset = Data.TensorDataset(test_data, test_label)
    test_loader = Data.DataLoader(
        dataset=test_dataset,
        batch_size=32,
        shuffle=True,
    )

    y_true = []
    y_pred = []
    for stp, (test_x, test_y) in enumerate(test_loader):
        test_output = net(test_x)
        _, pred_y = torch.max(test_output, 1)
        y_true.extend(test_y)
        y_pred.extend(pred_y)
    y_true = torch.tensor(y_true, device='cpu')
    y_pred = torch.tensor(y_pred, device='cpu')
    print("Accuracy:", accuracy_score(y_true, y_pred))
    print("Precision_score:", precision_score(y_true, y_pred, average='macro'))
    print("Recall_score:", recall_score(y_true, y_pred, average='macro'))
    print("F1_score", f1_score(y_true, y_pred, average='macro'))


def predict(model, file):
    spect = get_features(file)
    data = torch.Tensor(spect)
    data = data.unsqueeze(0)
    data = data.unsqueeze(0)
    data = data.to(device)

    output = model(data)
    confidence, pred_y = torch.max(output, 1)
    print("识别结果为：", pred_y.cpu().numpy())


if __name__ == '__main__':
    data = np.load("data.npy")
    label = np.load("label.npy")

    #cnn = ResNet(ResidualBlock, [2, 2, 2]).to(device)
    #cnn = CNN().to(device)
    cnn = MLP().to(device)
    cnn = train_model(cnn, data, label, lr=0.01, batch_size=512, epoch=300)
    save_model(cnn, "mlp.pt")

    test_model(cnn, data, label)

    np.save("dropout_train_losses.npy", train_losses)
    np.save("dropout_train_accuracy.npy", train_accuracy)

    plt.plot(train_losses, label='Training loss')
    plt.plot(train_accuracy, label='Accuracy')
    plt.xlabel("epoch")
    plt.ylabel("loss/accuracy")
    plt.legend()
    plt.show()