Pytorch同时迭代两个数据集

╰+攻爆jí腚メ 2022-11-11 05:50 263阅读 0赞

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