mirror of https://github.com/hpcaitech/ColossalAI
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from torch.utils.data import DataLoader
|
|
|
|
from colossalai.builder import build_dataset
|
|
from colossalai.context import Config
|
|
|
|
train_data = dict(
|
|
dataset=dict(
|
|
type='CIFAR10Dataset',
|
|
root=Path(os.environ['DATA']),
|
|
train=True,
|
|
download=True,
|
|
transform_pipeline=[
|
|
dict(type='ToTensor'),
|
|
dict(type='Normalize',
|
|
mean=(0.5, 0.5, 0.5),
|
|
std=(0.5, 0.5, 0.5))
|
|
]),
|
|
dataloader=dict(batch_size=4, shuffle=True, num_workers=2)
|
|
)
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_cifar10_dataset():
|
|
global train_data
|
|
config = Config(train_data)
|
|
dataset = build_dataset(config.dataset)
|
|
dataloader = DataLoader(dataset=dataset, **config.dataloader)
|
|
data_iter = iter(dataloader)
|
|
img, label = data_iter.next()
|
|
|
|
assert isinstance(img, list) and isinstance(label, list), \
|
|
f'expected the img and label to be list but got {type(img)} and {type(label)}'
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_cifar10_dataset()
|