基于CNN的多分类模型Python代码解析

概述

本代码实现了一个基于卷积神经网络（CNN）的多分类模型，并使用交叉熵损失函数进行训练。该模型应用于一个包含六个类别的分类问题，并使用准确率、NAR（错误分类率）和混淆矩阵等指标评估其性能。

代码解析

1. 日志记录

class Logger(object):
    def __init__(self, filename='default.log', stream=sys.stdout):
        self.terminal = stream
        self.log = open(filename, 'w')

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass

sys.stdout = Logger('result.log', sys.stdout)

这段代码定义了一个Logger类，用于将程序输出同时打印到控制台和日志文件result.log中。

2. 导入库

import argparse
import numpy as np
import time
from sklearn.metrics import accuracy_score, confusion_matrix
from net.cnn_gmlp_model import CNN
from mydataset import MyDataset
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from tqdm import *

这段代码导入了程序所需的库，包括用于参数解析的argparse，用于数值计算的numpy，用于时间统计的time，用于评估模型性能的sklearn.metrics，用于定义CNN模型的net.cnn_gmlp_model，用于加载数据的mydataset，用于深度学习的torch，用于数据可视化的matplotlib.pyplot和seaborn，以及用于进度条显示的tqdm。

3. 测试函数

def test(model, dataset, criterion):
    model.eval()
    total_batch_num = 0.
    val_loss = 0
    prediction = []
    labels = []
    feature_list = torch.tensor([])
    for (step, i) in enumerate(dataset):
        total_batch_num = total_batch_num + 1
        batch_x = i['data']
        batch_y = i['label']
        batch_x = torch.unsqueeze(batch_x, dim=1)
        batch_x = batch_x.float()
        if torch.cuda.is_available():
            batch_x = batch_x.cuda()
            batch_y = batch_y.cuda()
        feature, probs = model(batch_x)
        feature = feature.cpu().detach()
        batch_label = batch_y.unsqueeze(1).float()
        batch_label = batch_label.cpu().detach()
        feature_label = torch.cat((feature, batch_label), dim=1)
        feature_list = torch.cat((feature_list, feature_label), dim=0)
        loss = criterion(probs, batch_y)
        _, pred = torch.max(probs, dim=1)
        predi = pred.tolist()
        label = batch_y.tolist()
        val_loss += loss.item()
        prediction.extend(predi)
        labels.extend(label)
    accuracy = accuracy_score(labels, prediction)
    C = confusion_matrix(labels, prediction)
    return accuracy, val_loss / total_batch_num, feature_list, C

test函数用于评估模型在测试集上的性能。它接受模型、测试集和损失函数作为输入，并返回准确率、平均损失、特征列表和混淆矩阵。

4. 训练函数

def train(model, train_x, train_y, optimizer, criterion):
    model.train()
    model.zero_grad()
    _, probs = model(train_x)
    loss = criterion(probs, train_y)
    _, pred = torch.max(probs, dim=1)
    labels = train_y.tolist()
    predi = pred.tolist()
    loss.backward()
    optimizer.step()
    return labels, predi, loss.item()

train函数用于训练模型。它接受模型、训练数据、优化器和损失函数作为输入，并返回训练标签、预测结果和损失值。

5. 可视化函数

def draw(train_acc, train_loss, test_acc, test_loss):
    # 绘制训练集和测试集上的准确率和损失曲线
    ...

def draw_result(C):
    # 绘制混淆矩阵
    ...

draw函数用于绘制训练集和测试集上的准确率和损失曲线，而draw_result函数用于绘制混淆矩阵。

6. 主函数

def main(args):
    # 数据加载
    ...

    # 模型定义
    cnn_gmlp_model = {'CNN': CNN}
    model = cnn_gmlp_model[args.model]()
    if torch.cuda.is_available():
        model = model.cuda()

    # 优化器和损失函数
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
    criterion = nn.CrossEntropyLoss()

    # 模型训练和评估
    ...

if __name__ == '__main__':
    # 参数解析
    ...
    main(my_args)

main函数是程序的入口点。它首先加载数据，然后定义模型、优化器和损失函数。接下来，它训练模型并定期评估其性能。最后，它保存训练好的模型并可视化训练过程。

总结

本代码提供了一个完整的CNN多分类模型实现示例，包括数据加载、模型定义、训练、评估和结果可视化。该代码结构清晰，注释详细，方便读者理解和学习。