pytorch加载自己的数据集进阶

import torchfrom torch.autograd import Variablefrom torchvision import transformsfrom torch.utils.data import Dataset, DataLoaderfrom PIL import Image

root = "/home/zlab/zhangshun/torch1/data_et/"# -----------------ready the dataset--------------------------def default_loader(path):
    return Image.open(path).convert('RGB')
    class MyDataset (Dataset):
    # 构造函数带有默认参数
    def __init__(self, txt, transform=None, target_transform=None, loader=default_loader):
        fh = open(txt, 'r')
        imgs = []
        for line in fh:
            # 移除字符串首尾的换行符
            # 删除末尾空
            # 以空格为分隔符 将字符串分成
            line = line.strip('\n')
            line = line.rstrip()
            words = line.split()
            imgs.append((words[0], int(words[1])))#imgs中包含有图像路径和标签
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader    def __getitem__(self, index):
        fn, label = self.imgs[index]
        #调用定义的loader方法
        img = self.loader(fn)
        if self.transform is not None:
            img = self.transform(img)
        return img, label    def __len__(self):
        return len(self.imgs)train_data = MyDataset(txt=root + 'train.txt', transform=transforms.ToTensor())test_data = MyDataset(txt=root + 'test.txt', transform=transforms.ToTensor())#train_data 和test_data包含多有的训练与测试数据，调用DataLoader批量加载train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True)test_loader = DataLoader(dataset=test_data, batch_size=64)

from torchvision import transformsfrom torch.utils.data import Dataset, DataLoaderfrom PIL import Image
root = "/home/zlab/zhangshun/torch1/data_et/"# -----------------ready the dataset--------------------------def default_loader(path):
    return Image.open(path).convert('RGB')class MyDataset(Dataset):
    # 构造函数带有默认参数
    def __init__(self, txt, transform=None, target_transform=None, loader=default_loader):
        fh = open(txt, 'r')
        imgs = []
        for line in fh:
            line = line.strip('\n')
            line = line.rstrip()
            words = line.split()
            imgs.append((words[0], int(words[1])))  # imgs中包含有图像路径和标签
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader    def __getitem__(self, index):
        fn, label = self.imgs[index]
        # 调用定义的loader方法
        img = self.loader(fn)
        if self.transform is not None:
            img = self.transform(img)
        return img, label    def __len__(self):
        return len(self.imgs)
    
    transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.CenterCrop(224),
            transforms.ToTensor()])train_data = MyDataset(txt=root + 'train.txt', transform=transform)test_data = MyDataset(txt=root + 'test.txt', transform=transform)# train_data 和test_data包含多有的训练与测试数据，调用DataLoader批量加载train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True)test_loader = DataLoader(dataset=test_data, batch_size=64)

使用ImageFolder

import osfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms
TRAIN_DIR = 'train'VALIDATION_DIR = 'valid'MEAN_RGB = (0.485, 0.456, 0.406)VAR_RGB = (0.229, 0.224, 0.225)transform_train = transforms.Compose([
    transforms.RandomSizedCrop(224, scale=(0.2, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(MEAN_RGB, VAR_RGB),])transform_test = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(MEAN_RGB, VAR_RGB),])def get_imagenet_dataset(batch_size, dataset_root='./dataset/imagenet/', dataset_tpye='train'):
    if dataset_tpye == 'train':
        train_dataset_root = os.path.join(dataset_root, TRAIN_DIR)
        trainset = datasets.ImageFolder(root=train_dataset_root, transform=transform_train)
        trainloader = DataLoader(trainset,
                                 batch_size=batch_size,
                                 shuffle=True,
                                 num_workers=8,
                                 pin_memory=True,
                                 drop_last=False)
        print('Succeed to init ImageNet train DataLoader!')
        return trainloader    elif dataset_tpye == 'val' or dataset_tpye == 'valid':
        val_dataset_root = os.path.join(dataset_root, VALIDATION_DIR)
        valset = datasets.ImageFolder(root=val_dataset_root, transform=transform_test)
        valloader = DataLoader(valset,
                               batch_size=batch_size,
                               shuffle=False,
                               num_workers=8,
                               pin_memory=False,
                               drop_last=False)
        print('Succeed to init ImageNet val DataLoader!')
        return valloader    else:
        raise Exception('IMAGENET DataLoader: Unknown dataset type -- %s' % dataset_tpye)

               

查看更多关于pytorch加载自己的数据集进阶的详细内容...