pytorch k 折

原创已于 2022-03-25 00:46:52 修改 · 2k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #神经网络 #python

于 2022-03-25 00:45:25 首次发布

本文介绍了一种使用PyTorch实现的图像分类任务流程，包括数据集划分、模型搭建、训练过程及验证评估等关键步骤。通过五折交叉验证提高模型泛化能力。

部署运行你感兴趣的模型镜像

### Load dataset

trainset = datasets.ImageFolder(train_dir, transform=transform)
testset = datasets.ImageFolder(test_dir, transform=transform)

import torch
from sklearn.model_selection import KFold

data_induce = np.arange(0, len(trainset))  # 将“训练集”分为训练集和验证集
kf = KFold(n_splits=5)  # 分成 5 份

for k, (train_index, val_index) in enumerate(kf.split(data_induce)):

    print('{} - FOLD'.format(k))

    train_subset = torch.utils.data.dataset.Subset(trainset, train_index)
    val_subset = torch.utils.data.dataset.Subset(trainset, val_index)
    trainloader = DataLoader(dataset=train_subset, batch_size=bs, pin_memory=True)
    valloader = DataLoader(dataset=val_subset, batch_size=bs, pin_memory=True)

    ### Build model
    criterion = nn.CrossEntropyLoss()
    model = torchvision.models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)  ## 2 classes
    model.to(DEVICE)

    ### Design optimizer
    optimizer = optim.Adam(model.parameters(), lr=modellr)

    for epoch in range(1, EPOCHS + 1):
        adjust_learning_rate(optimizer, epoch)
        train(model, DEVICE, trainloader, optimizer, epoch, k)
        val(model, DEVICE, valloader, k)
    torch.save(model, 'model_{}_fold.pth'.format(k))

在原创基础上加工，原创有点找不到了，仅作记录

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch

Cuda

PyTorch 是一个开源的 Python 机器学习库，基于 Torch 库，底层由 C++ 实现，应用于人工智能领域，如计算机视觉和自然语言处理