Pytorch同时迭代两个数据集

from __future__ import print_function, division, absolute_import
import argparse
import os
import shutil
import time

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets

import sys

sys.path.append('.')
import pretrainedmodels
import pretrainedmodels.utils


scale = 0.875
print("=> using pre-trained parameters '{}'".format('imagenet'))
a1 = 0.4970
model_1 = pretrainedmodels.__dict__['resnet18'](num_classes=1000,pretrained='imagenet')
a2 = 0.5030
model_2 = pretrainedmodels.__dict__['vgg13'](num_classes=1000,pretrained='imagenet')

print('Images transformed from size {} to {}'.format(
    int(round(max(model_1.input_size) / scale)),
    model_1.input_size))

val_tf_1 = pretrainedmodels.utils.TransformImage(
    model_1,
    scale=scale,
    preserve_aspect_ratio=True
)

val_tf_2 = pretrainedmodels.utils.TransformImage(
    model_2,
    scale=scale,
    preserve_aspect_ratio=True
)

val_loader_1 = torch.utils.data.DataLoader(
    datasets.ImageFolder(valdir, val_tf_1),
    batch_size=20, shuffle=False,
    num_workers=4, pin_memory=True)

val_loader_2 = torch.utils.data.DataLoader(
    datasets.ImageFolder(valdir, val_tf_2),
    batch_size=20, shuffle=False,
    num_workers=4, pin_memory=True)
criterion = nn.CrossEntropyLoss().cuda()
model_1 = torch.nn.DataParallel(model_1).cuda()
model_2 = torch.nn.DataParallel(model_2).cuda()
validate_2(val_loader_1,val_loader_2, a1, model_1, a2, model_2, criterion)


def validate_2(val_loader_1, val_loader_2, a1, model_1, a2, model_2, criterion):
    with torch.no_grad():
        batch_time = AverageMeter()
        losses = AverageMeter()
        top1 = AverageMeter()
        top5 = AverageMeter()

        # switch to evaluate mode
        model_1.eval()
        model_2.eval()

        end = time.time()
        for i, data in enumerate(zip(val_loader_1, val_loader_2)):
            inputs1 = data[0][0].cuda();
            labels1 = data[0][1].cuda();
            inputs2 = data[1][0].cuda();
            labels2 = data[1][1].cuda();
            
            # compute output
            output_1 = model_1(inputs1)
            output_2 = model_2(inputs2)
            loss_1 = criterion(output_1, labels1)
            loss_2 = criterion(output_2, labels2)
        return output_1，output_2