import torch
import pandas as pd

data_path = '../sample_data/wikipedia_split_psgs_w100.tsv'

data = pd.read_csv(data_path, sep='\t')

data.head()

	id	text	title
0	1	지미 카터 제임스 얼 "지미" 카터 주니어(, 1924년 10월 1일 ~ )는 민주...	지미 카터
1	2	등용법을 내세웠다. 1976년 미합중국 (미국) 제39대 대통령 선거에 민주당 후보...	지미 카터
2	3	구출 실패를 이유로 1980년 대통령 선거에서 공화당의 로널드 레이건 후보에게 져 ...	지미 카터
3	4	그는 이에 대해 애매한 태도를 보였고, 이는 후에 대한민국 내에서 고조되는 반미 운...	지미 카터
4	5	카터는 카터 행정부 이후 미국이 북핵 위기, 코소보 전쟁, 이라크 전쟁과 같이 미국...	지미 카터

Single Process

Map-style Dataset

from torch.utils.data import Dataset

class MyMapDataset(Dataset):

    def __init__(self, data):
        self.data = data

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

    def __getitem__(self, index):
        return self.data['text'][index]

dataset = MyMapDataset(data)

dataset[0]

'지미 카터 제임스 얼 "지미" 카터 주니어(, 1924년 10월 1일 ~ )는 민주당 출신 미국 39대 대통령 (1977년 ~ 1981년)이다. 지미 카터는 조지아주 섬터 카운티 플레인스 마을에서 태어났다. 조지아 공과대학교를 졸업하였다. 그 후 해군에 들어가 전함·원자력·잠수함의 승무원으로 일하였다. 1953년 미국 해군 대위로 예편하였고 이후 땅콩·면화 등을 가꿔 많은 돈을 벌었다. 그의 별명이 "땅콩 농부" (Peanut Farmer)로 알려졌다. 1962년 조지아 주 상원 의원 선거에서 낙선하나 그 선거가 부정선거 였음을 입증하게 되어 당선되고, 1966년 조지아 주 지사 선거에 낙선하지만 1970년 조지아 주 지사를 역임했다. 대통령이 되기 전 조지아주 상원의원을 두번 연임했으며, 1971년부터 1975년까지 조지아 지사로 근무했다. 조지아 주지사로 지내면서, 미국에 사는 흑인'

len(dataset)

999

Iterable Dataset

from torch.utils.data import IterableDataset

class MyIterableDataset(IterableDataset):

    def __init__(self, data_path):
        self.data_path = data_path

    def __iter__(self):
        iter_csv = pd.read_csv(self.data_path, sep='\t', iterator=True, chunksize=1)
        for line in iter_csv:
            line = line['text'].item()
            yield line

idataset = MyIterableDataset(data_path)

for i in idataset:
    print(i)
    break

지미 카터 제임스 얼 "지미" 카터 주니어(, 1924년 10월 1일 ~ )는 민주당 출신 미국 39대 대통령 (1977년 ~ 1981년)이다. 지미 카터는 조지아주 섬터 카운티 플레인스 마을에서 태어났다. 조지아 공과대학교를 졸업하였다. 그 후 해군에 들어가 전함·원자력·잠수함의 승무원으로 일하였다. 1953년 미국 해군 대위로 예편하였고 이후 땅콩·면화 등을 가꿔 많은 돈을 벌었다. 그의 별명이 "땅콩 농부" (Peanut Farmer)로 알려졌다. 1962년 조지아 주 상원 의원 선거에서 낙선하나 그 선거가 부정선거 였음을 입증하게 되어 당선되고, 1966년 조지아 주 지사 선거에 낙선하지만 1970년 조지아 주 지사를 역임했다. 대통령이 되기 전 조지아주 상원의원을 두번 연임했으며, 1971년부터 1975년까지 조지아 지사로 근무했다. 조지아 주지사로 지내면서, 미국에 사는 흑인

Going Parallel

Map-style Dataset

from torch.utils.data import DataLoader, Dataset, IterableDataset
import time

class MyMapDataset(Dataset):
    
    def __init__(self, data):
        self.data = data
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        worker = torch.utils.data.get_worker_info()
        worker_id = worker.id if worker is not None else -1
        
        start = time.time()
        time.sleep(0.1)
        end = time.time()
        
        return self.data[index], worker_id, start, end

data = range(16)
map_dataset = MyMapDataset(data)

loader = DataLoader(map_dataset, batch_size=4, num_workers=0)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([-1, -1, -1, -1])
tensor([4, 5, 6, 7]) tensor([-1, -1, -1, -1])
tensor([ 8,  9, 10, 11]) tensor([-1, -1, -1, -1])
tensor([12, 13, 14, 15]) tensor([-1, -1, -1, -1])

loader = DataLoader(map_dataset, batch_size=4, num_workers=1)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([4, 5, 6, 7]) tensor([0, 0, 0, 0])
tensor([ 8,  9, 10, 11]) tensor([0, 0, 0, 0])
tensor([12, 13, 14, 15]) tensor([0, 0, 0, 0])

loader = DataLoader(map_dataset, batch_size=4, num_workers=2)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([4, 5, 6, 7]) tensor([1, 1, 1, 1])
tensor([ 8,  9, 10, 11]) tensor([0, 0, 0, 0])
tensor([12, 13, 14, 15]) tensor([1, 1, 1, 1])

loader = DataLoader(map_dataset, batch_size=4, num_workers=3)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([4, 5, 6, 7]) tensor([1, 1, 1, 1])
tensor([ 8,  9, 10, 11]) tensor([2, 2, 2, 2])
tensor([12, 13, 14, 15]) tensor([0, 0, 0, 0])

Iterable Dataset

from torch.utils.data import DataLoader, Dataset, IterableDataset
import time

class MyIterableDataset(IterableDataset):
    
    def __init__(self, data):
        self.data = data
    
    def __iter__(self):
        for x in self.data:
            worker = torch.utils.data.get_worker_info()
            worker_id = worker.id if worker is not None else -1
        
            start = time.time()
            time.sleep(0.1)
            end = time.time()
        
            yield x, worker_id, start, end

iterable_dataset = MyIterableDataset(data)

loader = DataLoader(iterable_dataset, batch_size=4, num_workers=0)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([-1, -1, -1, -1])
tensor([4, 5, 6, 7]) tensor([-1, -1, -1, -1])
tensor([ 8,  9, 10, 11]) tensor([-1, -1, -1, -1])
tensor([12, 13, 14, 15]) tensor([-1, -1, -1, -1])

loader = DataLoader(iterable_dataset, batch_size=4, num_workers=1)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([4, 5, 6, 7]) tensor([0, 0, 0, 0])
tensor([ 8,  9, 10, 11]) tensor([0, 0, 0, 0])
tensor([12, 13, 14, 15]) tensor([0, 0, 0, 0])

loader = DataLoader(iterable_dataset, batch_size=4, num_workers=2)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([0, 1, 2, 3]) tensor([1, 1, 1, 1])
tensor([4, 5, 6, 7]) tensor([0, 0, 0, 0])
tensor([4, 5, 6, 7]) tensor([1, 1, 1, 1])
tensor([ 8,  9, 10, 11]) tensor([0, 0, 0, 0])
tensor([ 8,  9, 10, 11]) tensor([1, 1, 1, 1])
tensor([12, 13, 14, 15]) tensor([0, 0, 0, 0])
tensor([12, 13, 14, 15]) tensor([1, 1, 1, 1])

loader = DataLoader(iterable_dataset, batch_size=4, num_workers=3)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([0, 1, 2, 3]) tensor([1, 1, 1, 1])
tensor([0, 1, 2, 3]) tensor([2, 2, 2, 2])
tensor([4, 5, 6, 7]) tensor([0, 0, 0, 0])
tensor([4, 5, 6, 7]) tensor([1, 1, 1, 1])
tensor([4, 5, 6, 7]) tensor([2, 2, 2, 2])
tensor([ 8,  9, 10, 11]) tensor([0, 0, 0, 0])
tensor([ 8,  9, 10, 11]) tensor([1, 1, 1, 1])
tensor([ 8,  9, 10, 11]) tensor([2, 2, 2, 2])
tensor([12, 13, 14, 15]) tensor([0, 0, 0, 0])
tensor([12, 13, 14, 15]) tensor([1, 1, 1, 1])
tensor([12, 13, 14, 15]) tensor([2, 2, 2, 2])

def worker_init_fn(_):
    worker_info = torch.utils.data.get_worker_info()
    
    dataset = worker_info.dataset
    worker_id = worker_info.id
    split_size = len(dataset.data) // worker_info.num_workers
    
    dataset.data = dataset.data[worker_id * split_size: (worker_id + 1) * split_size]

loader = DataLoader(iterable_dataset, batch_size=4, num_workers=2, worker_init_fn=worker_init_fn)
for d in loader:
    batch, worker_ids, starts, ends = d
    print(batch, worker_ids)

tensor([0, 1, 2, 3]) tensor([0, 0, 0, 0])
tensor([ 8,  9, 10, 11]) tensor([1, 1, 1, 1])
tensor([4, 5, 6, 7]) tensor([0, 0, 0, 0])
tensor([12, 13, 14, 15]) tensor([1, 1, 1, 1])