To fetch a item from generator, next() can be used: Return the next item from the iterator.
If a variable is not a generator, next() can be used along with iter().
Python code snippet:1
2
3
4
5
6
7
8
9
10
11
12
13a=[1,2,3]
next(a)
# output: TypeError: 'list' object is not an iterator
b=iter(a)
next(b)
# output: 1
next(b)
# output: 2
next(b)
# output: 3
next(b)
# StopIteration
Example:1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26import torch
import torchvision
import torchvision.transforms as transforms
batch_size = 256
# dataset construction
transform = transforms.Compose([
transforms.ToTensor(), # convert to tensor
transforms.Lambda(lambda x: x.view(image_dim)) # flatten into vector
])
train_set = torchvision.datasets.FashionMNIST(
root='./data/FashionMNIST'
,train=True
,download=True
,transform=transform
)
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=batch_size
)
# Fetch images by next() function
# Since the obj returned by DataLoader was not iterator, I also used iter()
images = next(iter(train_loader))