Function eval_unet

examples/mlperf/model_eval.py:362–481

(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
                inception:FidInceptionV3, clip:OpenClipEncoder)

  @Tensor.train(mode=False)
  def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
                inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
    # Eval is divided into 5 jits, one per model
    # It doesn't make sense to merge these jits, e.g. unet repeats 50 times in isolation; images fork to separate inception/clip
    # We're generating and scoring 30,000 images per eval, and all the data can flow through one jit at a time
    # To maximize throughput for each jit, we have only one model/jit on the GPU at a time, and pool outputs from each jit off-GPU
    for model in (unet, first_stage, inception, clip):
      Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])

    uc_written = False
    models = (cond_stage, unet, first_stage, inception, clip)
    jits = (jit_context:=TinyJit(cond_stage.embed_tokens), denoise_step, vae_decode, jit_inception:=TinyJit(inception),
            jit_clip:=TinyJit(clip.get_clip_score))
    all_bs = (CONTEXT_BS, DENOISE_BS, DECODE_BS, INCEPTION_BS, CLIP_BS)
    if (EVAL_SAMPLES:=getenv("EVAL_SAMPLES", 0)) and EVAL_SAMPLES > 0:
      eval_inputs = eval_inputs[0:EVAL_SAMPLES]
    output_shapes = [(ns:=len(eval_inputs),77), (ns,77,1024), (ns,4,64,64), (ns,3,512,512), (ns,2048), (ns,)]
    # Writing progress to disk lets us resume eval if we crash
    stages = ["tokens", "embeds", "latents", "imgs", "inception", "clip"]
    disk_tensor_names, disk_tensor_shapes = stages + ["end", "uc"], output_shapes + [(6,), (1,77,1024)]
    if not all(os.path.exists(f"{EVAL_CKPT_DIR}/{name}.bytes") for name in disk_tensor_names):
      for name, shape in zip(disk_tensor_names, disk_tensor_shapes):
        file = Path(f"{EVAL_CKPT_DIR}/{name}.bytes")
        file.unlink(missing_ok=True)
        with file.open("wb") as f: f.truncate(prod(shape) * 4)
    progress = {name: Tensor.empty(*shape, device=f"disk:{EVAL_CKPT_DIR}/{name}.bytes", dtype=dtypes.int if name in {"tokens", "end"} else dtypes.float)
                for name, shape in zip(disk_tensor_names, disk_tensor_shapes)}

    def embed_tokens(tokens:Tensor) -> Tensor:
      nonlocal uc_written
      if not uc_written:
        with Context(BEAM=0): progress["uc"].assign(cond_stage.embed_tokens(cond_stage.tokenize("").to(GPUS)).to("CPU").realize()).realize()
        uc_written = True
      return jit_context(shard_tensor(tokens))

    def generate_latents(embeds:Tensor) -> Tensor:
      uc_c = Tensor.stack(progress["uc"].to("CPU").expand(bs, 77, 1024), embeds, dim=1).reshape(-1, 77, 1024)
      uc_c = shard_tensor(uc_c)
      x = shard_tensor(Tensor.randn(bs,4,64,64))
      for step_idx, timestep in enumerate(tqdm(eval_timesteps)):
        reversed_idx = Tensor([50 - step_idx - 1], device=GPUS)
        alpha_prev = eval_alphas_prev[reversed_idx]
        ts = Tensor.full(bs, fill_value=timestep, dtype=dtypes.int, device="CPU")
        ts_ts = shard_tensor(ts.cat(ts))
        ts = shard_tensor(ts)
        sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
        sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
        x_x = shard_tensor(Tensor.stack(x.to("CPU"), x.to("CPU"), dim=1).reshape(-1, 4, 64, 64))
        x.assign(denoise_step(x, x_x, ts_ts, uc_c, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t, alpha_prev, unet, GPUS)).realize()
      return x

    def decode_latents(latents:Tensor) -> Tensor: return vae_decode(shard_tensor(latents), first_stage, disable_beam=True)
    def generate_inception(imgs:Tensor) -> Tensor: return jit_inception(shard_tensor(imgs))[:,:,0,0]

    def calc_clip_scores(batch:Tensor, batch_tokens:Tensor) -> Tensor:
      # Tensor.interpolate does not yet support bicubic, so we use PIL
      batch = (batch.to(GPUS[0]).permute(0,2,3,1) * 255).clip(0, 255).cast(dtypes.uint8).numpy()
      batch = [np.array(PIL.Image.fromarray(batch[i]).resize((224,224), PIL.Image.BICUBIC)) for i in range(bs)]