Use all memory when compute a batch (train) is slower than webdataset to generate a new batch about webdataset HOT 3 CLOSED

batch_size = 256
train_paths = ["gs://bla/bla/bla1", "gs://bla/bla/bla2"]
val_paths = ["gs://bla/bla/bla1", "gs://bla/bla/bla2"]
workers = 4

def decode_image(element):
    img = io.BytesIO(element)
    img = Image.open(img)
    return np.array(img)


def parse(input):
    data = json.loads(input['json'])
    return input['jpg'], data['subfamilia_idx']


def norm_imgs(imgs):
    imgs[:, 0] = (imgs[:, 0] - 0.485) / 0.229
    imgs[:, 1] = (imgs[:, 1] - 0.456) / 0.224
    imgs[:, 2] = (imgs[:, 2] - 0.406) / 0.225
    return imgs


augmentation = build_aug()
def decode_batch(batch):
    x, y = batch
    x = np.stack([decode_image(img) for img in x], 0)
    x = augmentation(images=x)
    x = np.rollaxis(x, 3, 1).astype('float32')
    x = norm_imgs(x)
    x = torch.from_numpy(x)
    return (x, torch.Tensor(np.array(y, dtype=np.int64)).type(torch.int64))


def decode_batch_val(batch):
    x, y = batch
    x = np.stack([decode_image(img) for img in x], 0)
    x = np.rollaxis(x, 3, 1).astype('float32')
    x = norm_imgs(x)
    x = torch.from_numpy(x)
    return (x, torch.Tensor(np.array(y, dtype=np.int64)).type(torch.int64))

# the strategy to execute ao process are much faster than process one by one.
dataset = wds.Dataset(train_paths).map(parse).shuffle(1000).batched(batch_size).map(decode_batch)
dataset_val = wds.Dataset(val_paths).map(parse).batched(batch_size).map(decode_batch_val)

dataloader = wds.MultiDataset(dataset, workers=workers)
dataloader_val = wds.MultiDataset(dataset_val, workers=workers)

for epoch in range(1000):
  for i, (input_var, output_var) in enumerate(dataloader):
    x = input_var.to(device)
    y = output_var.to(device, non_blocking=True)
    # do training

  for i, (input_var, output_var) in enumerate(dataloader_val):
    x = input_var.to(device)
    y = output_var.to(device, non_blocking=True)
    # do eval