mirror of https://github.com/hpcaitech/ColossalAI
29 lines
741 B
Python
29 lines
741 B
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from torchvision import transforms, datasets
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_cifar10_dataset():
|
|
# build transform
|
|
transform_pipeline = [transforms.ToTensor()]
|
|
transform_pipeline = transforms.Compose(transform_pipeline)
|
|
|
|
# build dataset
|
|
dataset = datasets.CIFAR10(root=Path(os.environ['DATA']), train=True, download=True, transform=transform_pipeline)
|
|
|
|
# build dataloader
|
|
dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=2)
|
|
data_iter = iter(dataloader)
|
|
img, label = data_iter.next()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_cifar10_dataset()
|