（二）PyTorch学习笔记——pytorch构造数据集

Dataset类

torch.utils.data.Dataset 是一个表示数据集的抽象类. 你自己的数据集一般应该继承Dataset, 并且重写下面的方法:

__len__使用len(dataset) 可以返回数据集的大小
__getitem__ 支持索引, 以便于使用 dataset[i] 可以 获取第:math:i个样本(0索引)
数据集创建一个数据集类. 我们使用 __init__方法来读取csv文件, 使用 __getitem__读取图片. 这样可以使内存高效利用, 因为我们并不需要在内存中一次存储所有图片, 而是按照需要读取.

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader

# 参数和数据加载
input_size = 5
output_size = 2

batch_size = 4   #批量大小
data_size = 100   #数据集的数量

#自定义数据集（需要覆盖以下3个方法）
class RandomDataset(Dataset):
    # 自定义数据（数据集大小length，每个数据为size的向量），格式为torch
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)
    #返回一个数据（返回值可以改写，这里是(index,self.data[index]),可以改成（self.data[index]））
    #参数index固定，代表当前数据的index
    def __getitem__(self, index):
        print("__getitem__")
        return index,self.data[index]

    # 返回数据集长度
    def __len__(self):
        return self.len

#DataLoader把数据集包装，变成一个迭代器，每次取出一个值是长度为batch_size的小数据集，shuffle=True
#代表把数据打乱
rand_loader = DataLoader(dataset=RandomDataset(input_size, 100),
                         batch_size=batch_size, shuffle=True)

class Model(nn.Module):
    # Our model

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("  In Model: input size", input.size(),
              "output size", output.size())

        return output

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1: #gpu个数大于1，并行处理
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  model = nn.DataParallel(model)

if torch.cuda.is_available(): #支持gpu
   model.cuda()

for index,data in rand_loader:  #这里的index,data代表batch_size长度的数据，对应到方法def __getitem__(self, index):
    if torch.cuda.is_available():
        input_var = Variable(data.cuda())
    else:
        input_var = Variable(data)

    output = model(input_var)
    print("Outside: input size", input_var.size(),
          "output_size", output.size())