mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 01:17:28 +08:00
f45a1ab83c
* update a few examples to use compile * update mnist * add compile to vae and rename some stuff for simplicity * update reqs * use state in eval * GCN example with RNG + dropout * add a bit of prefetching
31 lines
844 B
Python
31 lines
844 B
Python
import numpy as np
|
|
from mlx.data.datasets import load_cifar10
|
|
|
|
|
|
def get_cifar10(batch_size, root=None):
|
|
tr = load_cifar10(root=root)
|
|
|
|
mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, 3))
|
|
std = np.array([0.229, 0.224, 0.225]).reshape((1, 1, 3))
|
|
|
|
def normalize(x):
|
|
x = x.astype("float32") / 255.0
|
|
return (x - mean) / std
|
|
|
|
tr_iter = (
|
|
tr.shuffle()
|
|
.to_stream()
|
|
.image_random_h_flip("image", prob=0.5)
|
|
.pad("image", 0, 4, 4, 0.0)
|
|
.pad("image", 1, 4, 4, 0.0)
|
|
.image_random_crop("image", 32, 32)
|
|
.key_transform("image", normalize)
|
|
.batch(batch_size)
|
|
.prefetch(4, 4)
|
|
)
|
|
|
|
test = load_cifar10(root=root, train=False)
|
|
test_iter = test.to_stream().key_transform("image", normalize).batch(batch_size)
|
|
|
|
return tr_iter, test_iter
|