pytorch加载自己的数据集源码分享

x = torch.linspace(1, 10, 10)   # 训练数据 linspace返回一个一维的张量,（最小值，最大值，多少个数）
print(x)
y = torch.linspace(10, 1, 10)   # 标签
print(y)

import torch.utils.data as Data
torch_dataset = Data.TensorDataset(x, y)  # 对给定的 tensor 数据，将他们包装成 dataset
#输出的结果是<torch.utils.data.dataset.TensorDataset object at 0x00000145BD93F1C0>，需要使用加载器进行加载，才能迭代遍历
print(torch_dataset)

BATCH_SIZE=5
loader = Data.DataLoader(#使用支持的默认的数据集加载的方式
    # 从数据库中每次抽出batch size个样本
    dataset=torch_dataset,       # torch TensorDataset format   加载数据集
    batch_size=BATCH_SIZE,       # mini batch size 5
    shuffle=False,                # 要不要打乱数据 (打乱比较好)
    num_workers=2,               # 多线程来读数据
)
 
def show_batch():
    for epoch in range(3):
        for step, (batch_x, batch_y) in enumerate(loader): #加载数据集的时候起的作用很奇怪
            # training
            print("steop:{}, batch_x:{}, batch_y:{}".format(step, batch_x, batch_y))
            print("*"*100)
if __name__ == '__main__':
    show_batch()

# 自定义数据集类
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, *args):
        super().__init__()
        # 初始化数据集包含的数据和标签
        pass
        
    def __getitem__(self, index):
        # 根据索引index从文件中读取一个数据
        # 对数据预处理
        # 返回数据和对应标签
        pass
    
    def __len__(self):
        # 返回数据集的大小
        return len()

'''
生成训练集和测试集，保存在txt文件中
'''
import os
import random


train_ratio = 0.6


test_ratio = 1-train_ratio

rootdata = r"dataset"

#数组定义
train_list, test_list = [],[]
data_list = []

class_flag = -1
# 将绝对路径加载进数组中
for a,b,c in os.walk(rootdata):#os.walk可以获得根路径、文件夹以及文件，并会一直进行迭代遍历下去，直至只有文件才会结束
    print(a)
    for i in range(len(c)):
        data_list.append(os.path.join(a,c[i]))

    for i in range(0,int(len(c)*train_ratio)):
        train_data = os.path.join(a, c[i])+'\t'+str(class_flag)+'\n' #class_flag表示分类的类别
        train_list.append(train_data)

    for i in range(int(len(c) * train_ratio),len(c)):
        test_data = os.path.join(a, c[i]) + '\t' + str(class_flag)+'\n'
        test_list.append(test_data)

    class_flag += 1 

print(train_list)
# 将数组的内容打乱顺序
random.shuffle(train_list)
random.shuffle(test_list)

#分别将绝对路径对应的数组内容写进文本文件里
with open('train.txt','w',encoding='UTF-8') as f:
    for train_img in train_list:
        f.write(str(train_img))

with open('test.txt','w',encoding='UTF-8') as f:
    for test_img in test_list:
        f.write(test_img)

def __init__(self, txt_path, train_flag=True):
     #初始化图片对应的变量imgs_info以及一些相关变量
     self.imgs_info = self.get_images(txt_path) #imgs_info保存了图片以及标签
     self.train_flag = train_flag

     self.train_tf = transforms.Compose([#对训练集的图片进行预处理
             transforms.Resize(224),
             transforms.RandomHorizontalFlip(),
             transforms.RandomVerticalFlip(),
             transforms.ToTensor(),
             transform_BZ
         ])
     self.val_tf = transforms.Compose([#对测试集的图片进行预处理
             transforms.Resize(224),
             transforms.ToTensor(),
             transform_BZ
         ])

def __getitem__(self, index):
     img_path, label = self.imgs_info[index]
     #打开图片，并将RGBA转换为RGB，这里是通过PIL库打开图片的
     img = Image.open(img_path)
     img = img.convert('RGB')
     img = self.padding_black(img) #将图片添加上黑边的
     if self.train_flag: #选择是训练集还是测试集
         img = self.train_tf(img)
     else:
         img = self.val_tf(img)
     label = int(label)

     return img, label

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

if __name__=='__main__':
    batch_size = 16

    # # 给训练集和测试集分别创建一个数据集加载器
    train_data = LoadData("train.txt", True)
    valid_data = LoadData("test.txt", False)


    train_dataloader = DataLoader(dataset=train_data, num_workers=4, pin_memory=True, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(dataset=valid_data, num_workers=4, pin_memory=True, batch_size=batch_size)

    for X, y in test_dataloader:
        print("Shape of X [N, C, H, W]: ", X.shape)
        print("Shape of y: ", y.shape, y.dtype)
        break