给下面的代码补充一些辅助的print信息：比如一开始print一下参数设置，然后最后要保存的输出的txt文件中
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/mixup_upupmodel/other/default/1')  # 添加上传的目录到模块搜索路径

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

class Args:
    lr = 0.0002
    bs = 64
    es = 200
    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 = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)

    # 数据存储位置
    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()

        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()