clean up mlperf dataloader import (#5940)

use tinygrad tqdm for dataset, and PIL Image is only needed for resnet
This commit is contained in:
chenyu 2024-08-06 17:10:08 -04:00 committed by GitHub
parent 7b6496f2e6
commit 1dab75ae37
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 5 deletions

View File

@ -1,13 +1,12 @@
import os, random, pickle, queue
from typing import List
from pathlib import Path
import numpy as np
from PIL import Image
from tqdm import tqdm
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, prod, Context, round_up
from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm
### ResNet
class MyQueue:
@ -41,6 +40,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
signal.signal(signal.SIGINT, lambda _, __: exit(0))
from extra.datasets.imagenet import center_crop, preprocess_train
from PIL import Image
with Context(DEBUG=0):
while (_recv := q_in.get()) is not None: