TensorData和Dataloader

'''
Description: torch--TensorData
Autor: 365JHWZGo
Date: 2021-11-15 21:42:12
LastEditors: 365JHWZGo
LastEditTime: 2021-11-15 21:53:45
'''

import torch
import torch.utils.data as Data
import numpy as np

BATCH_SIZE = 2

#numpy
#方法一
x = np.linspace(1,10,10)[:,np.newaxis]
#方法二
x = np.expand_dims(np.linspace(1,10,10),1)
y = np.square(x)

#tensor
x = torch.unsqueeze(torch.linspace(1,10,10),dim=1)
y = torch.square(x)

dataset = Data.TensorDataset(torch.from_numpy(x),torch.from_numpy(y))
train_loader = torch.utils.data.DataLoader(
    batch_size=BATCH_SIZE,
    dataset=dataset,
    shuffle=True,
    num_workers=2
)
if __name__ == '__main__':
    for step,(batch_x,batch_y) in enumerate(train_loader):
        print('step:{0}\nbatch_x:{1}\nbatch_y:{2}\n'.format(step,batch_x,batch_y))

运行结果：

step:0
batch_x:tensor([[5.],[3.]], dtype=torch.float64)
batch_y:tensor([[25.],[ 9.]], dtype=torch.float64)

step:1
batch_x:tensor([[9.],[1.]], dtype=torch.float64)
batch_y:tensor([[81.],[ 1.]], dtype=torch.float64)

step:2
batch_x:tensor([[6.],[4.]], dtype=torch.float64)
batch_y:tensor([[36.],[16.]], dtype=torch.float64)

step:3
batch_x:tensor([[2.],[7.]], dtype=torch.float64)
batch_y:tensor([[ 4.],[49.]], dtype=torch.float64)

step:4
batch_x:tensor([[10.],[ 8.]], dtype=torch.float64)
batch_y:tensor([[100.],[ 64.]], dtype=torch.float64)