pytorch k 折

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

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