fix: pmap clip32

dev/inference/wandb-backend.ipynb

@@ -36,7 +36,8 @@
     "VQGAN_REPO, VQGAN_COMMIT_ID = 'dalle-mini/vqgan_imagenet_f16_16384', None\n",
     "normalize_text = True\n",
     "latest_only = False   # log only latest or all versions\n",
-    "suffix = '_1'           # mainly for duplicate inference runs with a deleted version"
+    "suffix = '_1'           # mainly for duplicate inference runs with a deleted version\n",
+    "add_clip_32 = False"
    ]
   },
   {
@@ -51,7 +52,8 @@
     "VQGAN_REPO, VQGAN_COMMIT_ID = 'dalle-mini/vqgan_imagenet_f16_16384', None\n",
     "normalize_text = False\n",
     "latest_only = True   # log only latest or all versions\n",
-    "suffix = '_2'           # mainly for duplicate inference runs with a deleted version"
+    "suffix = '_2'           # mainly for duplicate inference runs with a deleted version\n",
+    "add_clip_32 = True"
    ]
   },
   {
@@ -82,7 +84,12 @@
     "clip = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
     "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
     "clip_params = replicate(clip.params)\n",
-    "vqgan_params = replicate(vqgan.params)"
+    "vqgan_params = replicate(vqgan.params)\n",
+    "\n",
+    "if add_clip_32:\n",
+    "    clip32 = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "    processor32 = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "    clip32_params = replicate(clip32.params)"
    ]
   },
   {
@@ -98,8 +105,14 @@
     "\n",
     "@partial(jax.pmap, axis_name=\"batch\")\n",
     "def p_clip(inputs):\n",
-    "    logits = clip(**inputs).logits_per_image\n",
-    "    return logits"
+    "    logits = clip(params=clip_params, **inputs).logits_per_image\n",
+    "    return logits\n",
+    "\n",
+    "if add_clip_32:\n",
+    "    @partial(jax.pmap, axis_name=\"batch\")\n",
+    "    def p_clip32(inputs):\n",
+    "        logits = clip32(params=clip32_params, **inputs).logits_per_image\n",
+    "        return logits"
    ]
   },
   {
@@ -158,7 +171,7 @@
     "# retrieve inference run details\n",
     "def get_last_inference_version(run_id):\n",
     "    try:\n",
-    "        inference_run = api.run(f'dalle-mini/dalle-mini/inf-{run_id}{suffix}')\n",
+    "        inference_run = api.run(f'dalle-mini/dalle-mini/{run_id}-clip16{suffix}')\n",
     "        return inference_run.summary.get('version', None)\n",
     "    except:\n",
     "        return None"
@@ -215,6 +228,8 @@
     "    print(f'Processing artifact: {artifact.name}')\n",
     "    version = int(artifact.version[1:])\n",
     "    results = []\n",
+    "    if add_clip_32:\n",
+    "        results32 = []\n",
     "    columns = ['Caption'] + [f'Image {i+1}' for i in range(top_k)] + [f'Score {i+1}' for i in range(top_k)]\n",
     "    \n",
     "    if latest_only:\n",
@@ -232,7 +247,7 @@
     "\n",
     "    # start/resume corresponding run\n",
     "    if run is None:\n",
-    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'inf-{run_id}{suffix}', resume='allow')\n",
+    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'{run_id}-clip16{suffix}', resume='allow')\n",
     "\n",
     "    # work in temporary directory\n",
     "    with tempfile.TemporaryDirectory() as tmp:\n",
@@ -283,7 +298,6 @@
     "            logits = logits.reshape(-1, num_images)\n",
     "            top_scores = logits.argsort()[:, -top_k:][..., ::-1]\n",
     "            logits = jax.device_get(logits)\n",
-    "\n",
     "            # add to results table\n",
     "            for i, (idx, scores, sample) in enumerate(zip(top_scores, logits, batch)):\n",
     "                if sample == padding_item: continue\n",
@@ -291,11 +305,68 @@
     "                top_images = [wandb.Image(cur_images[x]) for x in idx]\n",
     "                top_scores = [scores[x] for x in idx]\n",
     "                results.append([sample] + top_images + top_scores)\n",
+    "                \n",
+    "            # get clip 32 scores - TODO: this should be refactored as it is same code as above\n",
+    "            if add_clip_32:\n",
+    "                print('Calculating CLIP 32 scores')\n",
+    "                clip_inputs = processor32(text=batch, images=images, return_tensors='np', padding='max_length', max_length=77, truncation=True).data\n",
+    "                # each shard will have one prompt, images need to be reorganized to be associated to the correct shard\n",
+    "                images_per_prompt_indices = np.asarray(range(0, len(images), batch_size))\n",
+    "                clip_inputs['pixel_values'] = jnp.concatenate(list(clip_inputs['pixel_values'][images_per_prompt_indices + i] for i in range(batch_size)))\n",
+    "                clip_inputs = shard(clip_inputs)\n",
+    "                logits = p_clip32(clip_inputs)\n",
+    "                logits = logits.reshape(-1, num_images)\n",
+    "                top_scores = logits.argsort()[:, -top_k:][..., ::-1]\n",
+    "                logits = jax.device_get(logits)\n",
+    "                # add to results table\n",
+    "                for i, (idx, scores, sample) in enumerate(zip(top_scores, logits, batch)):\n",
+    "                    if sample == padding_item: continue\n",
+    "                    cur_images = [images[x] for x in images_per_prompt_indices + i]\n",
+    "                    top_images = [wandb.Image(cur_images[x]) for x in idx]\n",
+    "                    top_scores = [scores[x] for x in idx]\n",
+    "                    results32.append([sample] + top_images + top_scores)\n",
     "\n",
     "    # log results\n",
     "    table = wandb.Table(columns=columns, data=results)\n",
     "    run.log({'Samples': table, 'version': version})\n",
-    "    wandb.finish()"
+    "    wandb.finish()\n",
+    "    \n",
+    "    if add_clip_32:        \n",
+    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'{run_id}-clip32{suffix}', resume='allow')\n",
+    "        table = wandb.Table(columns=columns, data=results32)\n",
+    "        run.log({'Samples': table, 'version': version})\n",
+    "        wandb.finish()\n",
+    "        run = None  # ensure we don't log on this run"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fdcd09d6-079c-461a-a81a-d9e650d3b099",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "p_clip32"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7d86ceee-c9ac-4860-abad-410cadd16c3c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "clip_inputs['attention_mask'].shape, clip_inputs['pixel_values'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fbba4858-da2d-4dd5-97b7-ce3ab4746f96",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "clip_inputs['input_ids'].shape"
    ]
   },
   {
@@ -314,12 +385,10 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e1c04761-1016-47e9-925c-3a9ec6fec95a",
+   "id": "a7a5fdf5-3c6e-421b-96a8-5115f730328c",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "wandb.finish()"
-   ]
+   "source": []
   }
  ],
  "metadata": {