使用Pytorch转换和自定义数据集时出错

这个问题主要涉及 __getitem__ 在a中 pytorch Dataset 我在源代码中看到了元组和dict。

我一直在跟踪 this tutorial 用于在代码中创建dataset类,如下所示 this tutorial 关于迁移学习。它具有以下数据集定义。

class FaceLandmarksDataset(Dataset):
"""Face Landmarks dataset."""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.landmarks_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.landmarks_frame)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir,
                                self.landmarks_frame.iloc[idx, 0])
        image = io.imread(img_name)
        landmarks = self.landmarks_frame.iloc[idx, 1:].as_matrix()
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample

如你所见, __获取项目__ 返回包含两个条目的词典。 在迁移学习教程中,进行以下调用以转换数据集:

    data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = 'hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

use_gpu = torch.cuda.is_available()

inputs, classes = next(iter(dataloaders['train']))

最后一行代码试图对自定义数据集中的示例运行转换,从而导致代码中出现错误。

“dict”对象没有属性“size”

但是,如果教程数据集实现正确,它是否应该正确地使用转换?我自己的混合实现如下:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
from torch.utils.data import *
from skimage import io, transform
plt.ion()


class NumsDataset(Dataset):
    """Face Landmarks dataset."""

    def __init__(self, root_dir, transform=None):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.docs = []
        for file in os.listdir(root_dir):
            #print(file)
            if file.endswith(".txt"):
                path = os.path.join(root_dir, file)
                with open(path, 'r') as f:
                    self.docs.append( (  file , list(f.read()) ) ) #tup containing file, image values pairs
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self): #returns number of images
        i = 0
        for j in self.docs:
            i += len(j[1])
        return i

    def len2(self): #returns number of batches
        return len(self.docs)

    def __getitem__(self, idx):
        idx1 = idx // self.len2()
        idx2 = idx % self.len2()
        imglabel = self.docs[idx1][0] #label with filename for batch error calculation later
        imgdir = os.path.join(self.root_dir, self.docs[idx1][0].strip(".txt"))
        img = None
        l = idx2

        for file in os.listdir(imgdir):
            file = os.path.join(imgdir, file)
            if(l == 0):
                img = io.imread(file)
            l -= 1
        sample = (img , imglabel)
        sample ={'image': img, 'label': imglabel}
        if self.transform:
            sample = self.transform(sample)

        return sample




data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
data_dir = "images"
image_datasets = {x: NumsDataset(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=5) 
              for x in ['train', 'val']}

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = ["one", "two", "four"]

use_gpu = torch.cuda.is_available()
# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))

目录结构:

images
     /train
        /file1
            *.jpg
        /file2...
            *.jpg
        file1.txt
        file2.txt...
     /val
        /file1
            *.jpg
        /file2...
            *.jpg
        file1.txt
        file2.txt...

我返回的样本格式是否不正确?