给我的代码加上一个早停
import sys
import torch
from torch import optim, nn
from torch.utils.data import DataLoader
import time
import os
import numpy as np
import matplotlib.pyplot as plt

from myDataset import MyDataset, mixup_collate_fn
from dataset import MyDataSet
from model_resnet import TransformerResNet1

sys.path.append('/kaggle/input/seipypy')  # 添加上传的目录到模块搜索路径

# 假设类别数量是 10，需确认与你数据的实际类别数保持一致
NUM_CLASSES = 10

class Args:
    lr = 0.0002
    bs = 64
    es = 100
    weight_decay = 1e-4
    class_num = NUM_CLASSES
    arch = 'MODEL1'

args = Args()

torch.manual_seed(1234)

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

def main():
    print("CUDA Available: ", torch.cuda.is_available())
    if torch.cuda.is_available():
        print("CUDA Device Name: ", torch.cuda.get_device_name(0))

    source_root = r'/kaggle/input/seisei/train'
    target_root = r'/kaggle/input/seisei/test'
    source_classes = list(range(NUM_CLASSES))  # 定义类别范围

    # 使用 MyDataSet 类进行数据加载
    source_set = MyDataSet(source_root, source_classes, range(500))
    target_set = MyDataSet(target_root, source_classes, range(400))

    print(f'class2label: {source_set.class2label}')
    print('source_set sample num:', len(source_set))
    print('target_set sample num:', len(target_set))

    # 定义 DataLoader，可能需要根据需要调整
    train_loader = DataLoader(
        source_set,
        batch_size=args.bs,
        shuffle=True,
        num_workers=0
    )

    test_loader = DataLoader(
        target_set,
        batch_size=args.bs,
        shuffle=False,
        num_workers=0
    )

    # 创建模型
    model = TransformerResNet1(num_classes=args.class_num).to(device)

    # 定义损失函数
    loss_fn = nn.CrossEntropyLoss()

    # 定义优化器
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

    # 学习率调度器
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode='min',  # 根据验证损失调整
        factor=0.5,  # 每次减少学习率的倍率
        patience=5,  # 在验证损失无改善的情况下等待5个epoch
        min_lr=1e-6  # 设置学习率下限
    )

    # 数据存储位置
    resultdatafolder = f'./result/{args.arch}/{args.class_num}'
    i = 0
    while os.path.exists(f'{resultdatafolder}_{i}'):
        i += 1
    resultdatafolder = f'{resultdatafolder}_{i}'
    os.makedirs(resultdatafolder)

    train_model(model, train_loader, test_loader, loss_fn, optimizer, scheduler, device, resultdatafolder)
    test_acc = test_model(model, test_loader, loss_fn, device, resultdatafolder)
    print(f"Final Test Accuracy: {test_acc:.2f}%")

def train_model(model, train_loader, test_loader, loss_func, optimizer, scheduler, device, resultdatafolder):
    start = time.time()
    train_acc_list = []
    train_loss_list = []
    val_acc_list = []
    val_loss_list = []

    best_acc = 0.0
    lowest_loss = float('inf')
    best_acc_epoch = 0
    lowest_loss_epoch = 0

    for epoch in range(args.es):
        print(f"Epoch {epoch + 1} / {args.es} lr={optimizer.param_groups[0]['lr']}")

        train_acc, train_loss = train(model, train_loader, loss_func, optimizer, device)
        val_acc, val_loss, _, _ = test(model, test_loader, loss_func, device)

          # 更新调度器
        scheduler.step(val_loss)

        # 获取当前学习率
        current_lr = optimizer.param_groups[0]['lr']
        print(f"Epoch {epoch+1}/{args.es}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, LR: {current_lr:.6f}")

        train_acc_list.append(train_acc)
        train_loss_list.append(train_loss)
        val_acc_list.append(val_acc)
        val_loss_list.append(val_loss)

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), f'{resultdatafolder}/best.model')
            best_acc_epoch = epoch
        if val_loss < lowest_loss:
            lowest_loss = val_loss
            lowest_loss_epoch = epoch

        time_elapsed = time.time() - start
        print(f"train_acc: {(100 * train_acc):.2f}% \t train_loss: {train_loss:.4f}")
        print(f"val_acc: {(100 * val_acc):.2f}% \t val_loss: {val_loss:.4f}")
        print('Train time elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print("\n")

    with open(os.path.join(resultdatafolder, 'acc.txt'), mode='w', newline='') as f:
        f.write(f'best_acc={best_acc}\tbest_acc_epoch={best_acc_epoch}\n')
        f.write(f'lowest_loss={lowest_loss}\tlowest_loss_epoch={lowest_loss_epoch}\n')

    plot_acc_and_loss(train_acc_list, train_loss_list, val_acc_list, val_loss_list, args.es, resultdatafolder)

def test_model(model, test_loader, loss_func, device, resultdatafolder):
    model.eval()
    test_acc, test_loss, true_label_list, pred_label_list = test(model, test_loader, loss_func, device)

    path = os.path.join(resultdatafolder, 'acc.txt')
    with open(path, mode='a', newline='') as f:
        f.write(f'test_acc={test_acc}\ttest_loss={test_loss}\n')

    return test_acc

def train(model, train_loader, loss_func, optimizer, device):
    model.train()
    total_correct = 0
    total_loss = 0.0
    dataset_size = len(train_loader.dataset)

    for features, labels in train_loader:
        features, labels = features.to(device), labels.to(device)
        optimizer.zero_grad()

        outputs = model(features)
        loss = loss_func(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * features.size(0)
        total_correct += (outputs.argmax(1) == labels).sum().item()

    avg_loss = total_loss / dataset_size
    avg_acc = total_correct / dataset_size
    return avg_acc, avg_loss

def test(model, test_loader, loss_func, device):
    model.eval()
    total_correct = 0
    total_loss = 0.0
    dataset_size = len(test_loader.dataset)
    true_label_list = []
    pred_label_list = []

    with torch.no_grad():
        for features, labels in test_loader:
            features, labels = features.to(device), labels.to(device)

            outputs = model(features)
            loss = loss_func(outputs, labels)

            total_loss += loss.item() * features.size(0)
            total_correct += (outputs.argmax(1) == labels).sum().item()

            true_label_list.extend(labels.cpu().numpy())
            pred_label_list.extend(outputs.argmax(1).cpu().numpy())

    avg_loss = total_loss / dataset_size
    avg_acc = total_correct / dataset_size
    return avg_acc, avg_loss, true_label_list, pred_label_list

def plot_acc_and_loss(train_acc_list, train_loss_list, val_acc_list, val_loss_list, epoch, resultdatafolder):
    epochlist = list(range(1, epoch + 1))
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(epochlist, train_acc_list, 'b-', label='Train Accuracy')
    plt.plot(epochlist, val_acc_list, 'r-', label='Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Train and Validation Accuracy')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochlist, train_loss_list, 'b-', label='Train Loss')
    plt.plot(epochlist, val_loss_list, 'r-', label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Train and Validation Loss')
    plt.legend()

    plt.tight_layout()
    plt.savefig(f'{resultdatafolder}/acc_loss.png')
    plt.show()

if __name__ == '__main__':
    main()