preprocess_clip_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cuda\n",
      "Attributes DataFrame columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n",
      "Attributes DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Sideburns  Smiling  \\\n",
      "0    -1     -1        -1        -1          -1  ...         -1        1   \n",
      "1    -1     -1        -1         1          -1  ...         -1        1   \n",
      "2    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "3    -1     -1        -1        -1          -1  ...         -1       -1   \n",
      "4    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "\n",
      "   Straight_Hair  Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0              1         -1                 1           -1                 1   \n",
      "1             -1         -1                -1           -1                -1   \n",
      "2             -1          1                -1           -1                -1   \n",
      "3              1         -1                 1           -1                 1   \n",
      "4             -1         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  \n",
      "0                -1               -1      1  \n",
      "1                -1               -1      1  \n",
      "2                -1               -1      1  \n",
      "3                 1               -1      1  \n",
      "4                -1               -1      1  \n",
      "\n",
      "[5 rows x 41 columns]\n",
      "'image_id' column found in attributes DataFrame.\n",
      "Partition DataFrame columns: ['image_id', 'partition']\n",
      "Partition DataFrame head:\n",
      "      image_id  partition\n",
      "0  000001.jpg          0\n",
      "1  000002.jpg          0\n",
      "2  000003.jpg          0\n",
      "3  000004.jpg          0\n",
      "4  000005.jpg          0\n",
      "Bounding Box DataFrame columns: ['image_id', 'x_1', 'y_1', 'width', 'height']\n",
      "Bounding Box DataFrame head:\n",
      "      image_id  x_1  y_1  width  height\n",
      "0  000001.jpg   95   71    226     313\n",
      "1  000002.jpg   72   94    221     306\n",
      "2  000003.jpg  216   59     91     126\n",
      "3  000004.jpg  622  257    564     781\n",
      "4  000005.jpg  236  109    120     166\n",
      "After merging with partition, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition']\n",
      "Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Smiling  Straight_Hair  \\\n",
      "0    -1     -1        -1        -1          -1  ...        1              1   \n",
      "1    -1     -1        -1         1          -1  ...        1             -1   \n",
      "2    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "3    -1     -1        -1        -1          -1  ...       -1              1   \n",
      "4    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "\n",
      "   Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0         -1                 1           -1                 1   \n",
      "1         -1                -1           -1                -1   \n",
      "2          1                -1           -1                -1   \n",
      "3         -1                 1           -1                 1   \n",
      "4         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  partition  \n",
      "0                -1               -1      1          0  \n",
      "1                -1               -1      1          0  \n",
      "2                -1               -1      1          0  \n",
      "3                 1               -1      1          0  \n",
      "4                -1               -1      1          0  \n",
      "\n",
      "[5 rows x 42 columns]\n",
      "After merging with bounding box, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition', 'x_1', 'y_1', 'width', 'height']\n",
      "Final Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Wearing_Hat  \\\n",
      "0    -1     -1        -1        -1          -1  ...           -1   \n",
      "1    -1     -1        -1         1          -1  ...           -1   \n",
      "2    -1     -1         1        -1          -1  ...           -1   \n",
      "3    -1     -1        -1        -1          -1  ...           -1   \n",
      "4    -1     -1         1        -1          -1  ...           -1   \n",
      "\n",
      "   Wearing_Lipstick  Wearing_Necklace  Wearing_Necktie  Young  partition  x_1  \\\n",
      "0                 1                -1               -1      1          0   95   \n",
      "1                -1                -1               -1      1          0   72   \n",
      "2                -1                -1               -1      1          0  216   \n",
      "3                 1                 1               -1      1          0  622   \n",
      "4                 1                -1               -1      1          0  236   \n",
      "\n",
      "   y_1  width  height  \n",
      "0   71    226     313  \n",
      "1   94    221     306  \n",
      "2   59     91     126  \n",
      "3  257    564     781  \n",
      "4  109    120     166  \n",
      "\n",
      "[5 rows x 46 columns]\n",
      "Total samples before filtering: 202599\n",
      "Total samples after filtering partition 0: 162770\n",
      "Attribute names: ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings:   0%|          | 0/318 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index: 0, Text Prompt: A person with Arched_Eyebrows, Attractive, Brown_Hair, Heavy_Makeup, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Pointy_Nose, Smiling, Straight_Hair.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 1, Text Prompt: A person with Bags_Under_Eyes, Big_Nose, Brown_Hair, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Smiling, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 2, Text Prompt: A person with Big_Lips, Blurry, Male, Narrow_Eyes, No_Beard, Pointy_Nose, Wavy_Hair, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 3, Text Prompt: A person with Attractive, No_Beard, Pointy_Nose, Straight_Hair, Wearing_Earrings, Wearing_Lipstick, Wearing_Necklace, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 4, Text Prompt: A person with Arched_Eyebrows, Attractive, Big_Lips, Heavy_Makeup, Narrow_Eyes, No_Beard, Pointy_Nose, Wearing_Lipstick, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings: 100%|██████████| 318/318 [02:15<00:00,  2.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "嵌入已保存为 'text_embeddings.pt' 和 'image_embeddings.pt'\n"
     ]
    }
   ],
   "source": [
    "# preprocess_clip_embeddings.py\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import clip\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "# 配置参数\n",
    "IMG_DIR = '/root/autodl-tmp/celeba_datasets/img_align_celeba/img_align_celeba'\n",
    "ATTR_PATH = '/root/autodl-tmp/celeba_datasets/list_attr_celeba.txt'\n",
    "BBOX_PATH = '/root/autodl-tmp/celeba_datasets/list_bbox_celeba.txt'\n",
    "PARTITION_PATH = '/root/autodl-tmp/celeba_datasets/list_eval_partition.txt'\n",
    "\n",
    "OUTPUT_DIR = './clip_embeddings'  # 输出目录\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "IMAGE_SIZE = 224  # CLIP ViT-B/32的默认图像尺寸\n",
    "BATCH_SIZE = 512  # 根据GPU显存调整\n",
    "NUM_WORKERS = 4  # 根据CPU核心数调整\n",
    "\n",
    "# 设备设置\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "# 图像预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.48145466, 0.4578275, 0.40821073], \n",
    "                         [0.26862954, 0.26130258, 0.27577711])  # CLIP的归一化参数\n",
    "])\n",
    "\n",
    "# 数据集类\n",
    "class CelebAEmbeddingsDataset(Dataset):\n",
    "    def __init__(self, img_dir, attr_path, bbox_path, partition_path, transform=None, partition=0):\n",
    "        self.img_dir = img_dir\n",
    "        self.transform = transform\n",
    "\n",
    "        # 读取属性文件\n",
    "        try:\n",
    "            attr_df = pd.read_csv(attr_path, sep=',', header=0)  # 正确设置 header=0\n",
    "        except Exception as e:\n",
    "            print(f\"Error reading {attr_path}: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"Attributes DataFrame columns:\", attr_df.columns.tolist())\n",
    "        print(\"Attributes DataFrame head:\\n\", attr_df.head())\n",
    "\n",
    "        # 确认 'image_id' 列存在\n",
    "        if 'image_id' not in attr_df.columns:\n",
    "            raise ValueError(\"属性文件中缺少 'image_id' 列。请检查文件格式。\")\n",
    "        else:\n",
    "            print(\"'image_id' column found in attributes DataFrame.\")\n",
    "\n",
    "        # 读取分区文件\n",
    "        try:\n",
    "            partition_df = pd.read_csv(partition_path, sep=',', header=0, names=['image_id', 'partition'])\n",
    "        except Exception as e:\n",
    "            print(f\"Error reading {partition_path}: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"Partition DataFrame columns:\", partition_df.columns.tolist())\n",
    "        print(\"Partition DataFrame head:\\n\", partition_df.head())\n",
    "\n",
    "        # 读取边界框文件\n",
    "        try:\n",
    "            bbox_df = pd.read_csv(bbox_path, sep=',', header=0, names=['image_id', 'x_1', 'y_1', 'width', 'height'])\n",
    "        except Exception as e:\n",
    "            print(f\"Error reading {bbox_path}: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"Bounding Box DataFrame columns:\", bbox_df.columns.tolist())\n",
    "        print(\"Bounding Box DataFrame head:\\n\", bbox_df.head())\n",
    "\n",
    "        # 合并属性和分区数据\n",
    "        try:\n",
    "            attr_df = attr_df.merge(partition_df, on='image_id')\n",
    "        except Exception as e:\n",
    "            print(f\"Error merging attributes and partitions: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"After merging with partition, columns:\", attr_df.columns.tolist())\n",
    "        print(\"Merged DataFrame head:\\n\", attr_df.head())\n",
    "\n",
    "        # 合并属性、分区和边界框数据\n",
    "        try:\n",
    "            attr_df = attr_df.merge(bbox_df, on='image_id')\n",
    "        except Exception as e:\n",
    "            print(f\"Error merging attributes, partitions and bounding boxes: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"After merging with bounding box, columns:\", attr_df.columns.tolist())\n",
    "        print(\"Final Merged DataFrame head:\\n\", attr_df.head())\n",
    "\n",
    "        # 筛选分区\n",
    "        print(f\"Total samples before filtering: {len(attr_df)}\")\n",
    "        self.attr_df = attr_df[attr_df['partition'] == partition].reset_index(drop=True)\n",
    "        print(f\"Total samples after filtering partition {partition}: {len(self.attr_df)}\")\n",
    "\n",
    "        if len(self.attr_df) == 0:\n",
    "            raise ValueError(f\"训练集中没有样本。请检查数据文件的读取和过滤逻辑。\")\n",
    "\n",
    "        # 获取属性名称\n",
    "        # 属性列是除了 'image_id', 'partition', 'x_1', 'y_1', 'width', 'height' 之外的所有列\n",
    "        self.attr_names = [col for col in attr_df.columns if col not in ['image_id', 'partition', 'x_1', 'y_1', 'width', 'height']]\n",
    "        print(\"Attribute names:\", self.attr_names)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.attr_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # 获取图像路径\n",
    "        img_name = self.attr_df.iloc[idx, self.attr_df.columns.get_loc('image_id')]\n",
    "        img_path = os.path.join(self.img_dir, img_name)\n",
    "        try:\n",
    "            image = Image.open(img_path).convert('RGB')\n",
    "        except Exception as e:\n",
    "            print(f\"Error opening image {img_path}: {e}\")\n",
    "            # 可以选择跳过该样本，或返回一个默认图像\n",
    "            # 这里选择抛出异常\n",
    "            raise\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        # 获取属性标签并生成文本描述\n",
    "        attrs = self.attr_df.iloc[idx][self.attr_names].values\n",
    "        attrs = (attrs + 1) // 2  # 将-1转为0，1保持1\n",
    "        attrs = attrs.astype(int)\n",
    "        active_attrs = [attr for attr, val in zip(self.attr_names, attrs) if val == 1]\n",
    "\n",
    "        # 手动截断属性数量，例如，最多10个属性，以避免超长文本\n",
    "        max_attrs = 10\n",
    "        active_attrs = active_attrs[:max_attrs]\n",
    "\n",
    "        if len(active_attrs) == 0:\n",
    "            text_prompt = \"A person with no notable attributes.\"\n",
    "        else:\n",
    "            text_description = ', '.join(active_attrs)\n",
    "            text_prompt = f\"A person with {text_description}.\"\n",
    "\n",
    "        # 打印前5个样本的文本描述和图像形状\n",
    "        if idx < 5:\n",
    "            print(f\"Index: {idx}, Text Prompt: {text_prompt}\")\n",
    "            print(f\"Image shape: {image.shape}\")\n",
    "\n",
    "        return text_prompt, image\n",
    "\n",
    "# 初始化 CLIP 模型\n",
    "clip_model, preprocess = clip.load(\"ViT-B/32\", device=device)\n",
    "clip_model.eval()\n",
    "\n",
    "# 创建数据集和数据加载器\n",
    "try:\n",
    "    dataset = CelebAEmbeddingsDataset(\n",
    "        img_dir=IMG_DIR,\n",
    "        attr_path=ATTR_PATH,\n",
    "        bbox_path=BBOX_PATH,\n",
    "        partition_path=PARTITION_PATH,\n",
    "        transform=transform,\n",
    "        partition=0  # 0: train\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(f\"Error initializing dataset: {e}\")\n",
    "    raise\n",
    "\n",
    "if len(dataset) == 0:\n",
    "    raise ValueError(\"训练集中没有样本。请检查数据文件的读取和过滤逻辑。\")\n",
    "\n",
    "dataloader = DataLoader(\n",
    "    dataset,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=False,\n",
    "    num_workers=NUM_WORKERS,\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "# 计算并保存嵌入\n",
    "text_embeddings = []\n",
    "image_embeddings = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for text_prompts, images in tqdm(dataloader, desc=\"Computing Embeddings\"):\n",
    "        # 生成文本嵌入\n",
    "        try:\n",
    "            text_tokens = clip.tokenize(text_prompts, truncate=True).to(device)\n",
    "        except Exception as e:\n",
    "            print(f\"Error tokenizing text: {e}\")\n",
    "            raise\n",
    "\n",
    "        try:\n",
    "            text_emb = clip_model.encode_text(text_tokens).cpu()\n",
    "        except Exception as e:\n",
    "            print(f\"Error encoding text: {e}\")\n",
    "            raise\n",
    "\n",
    "        text_embeddings.append(text_emb)\n",
    "\n",
    "        # 生成图像嵌入\n",
    "        try:\n",
    "            images = images.to(device)\n",
    "            image_emb = clip_model.encode_image(images).cpu()\n",
    "        except Exception as e:\n",
    "            print(f\"Error encoding images: {e}\")\n",
    "            raise\n",
    "\n",
    "        image_embeddings.append(image_emb)\n",
    "\n",
    "# 合并所有嵌入\n",
    "if len(text_embeddings) == 0 or len(image_embeddings) == 0:\n",
    "    raise RuntimeError(\"未计算到任何嵌入。请检查数据加载和嵌入计算步骤。\")\n",
    "\n",
    "text_embeddings = torch.cat(text_embeddings, dim=0)\n",
    "image_embeddings = torch.cat(image_embeddings, dim=0)\n",
    "\n",
    "# 保存嵌入到文件\n",
    "torch.save(text_embeddings, os.path.join(OUTPUT_DIR, 'text_embeddings.pt'))\n",
    "torch.save(image_embeddings, os.path.join(OUTPUT_DIR, 'image_embeddings.pt'))\n",
    "\n",
    "print(\"嵌入已保存为 'text_embeddings.pt' 和 'image_embeddings.pt'\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_4722/3045327233.py:3: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  text_emb_sample = torch.load('./clip_embeddings/text_embeddings.pt')\n",
      "/tmp/ipykernel_4722/3045327233.py:4: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  image_emb_sample = torch.load('./clip_embeddings/image_embeddings.pt')\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "文本嵌入样本形状: torch.Size([162770, 512])\n",
      "图像嵌入样本形状: torch.Size([162770, 512])\n",
      "第一个文本嵌入: tensor([ 4.1901e-02,  3.7549e-01, -2.2180e-01,  1.8335e-01,  8.2458e-02,\n",
      "        -2.4121e-01, -2.4866e-01,  7.4268e-01,  7.8796e-02,  4.9707e-01,\n",
      "         1.7554e-01, -4.1534e-02,  3.5938e-01,  7.9163e-02,  2.6709e-01,\n",
      "         1.6577e-01,  5.2295e-01,  1.6785e-01, -2.5879e-01,  2.2693e-01,\n",
      "         4.9487e-01,  1.0559e-01, -2.4988e-01, -5.7434e-02, -1.7017e-01,\n",
      "        -4.9408e-02, -2.0850e-01,  1.2494e-01, -3.4497e-01, -1.4478e-01,\n",
      "         2.2058e-01, -2.9736e-01, -1.3196e-01,  1.3403e-01,  2.2412e-01,\n",
      "         1.3123e-01,  3.3997e-02,  7.4951e-02, -1.2549e-01,  2.5024e-01,\n",
      "        -1.8457e-01, -9.9121e-02, -1.6321e-01,  8.6365e-02, -1.2854e-01,\n",
      "         5.7520e-01,  1.4392e-01, -7.4524e-02, -9.7961e-02,  1.4917e-01,\n",
      "         8.5388e-02, -2.1866e-02,  1.6492e-01, -3.1714e-01, -9.6130e-02,\n",
      "        -1.4075e-01,  2.3010e-01,  5.4504e-02, -3.0914e-02,  1.6650e-01,\n",
      "         5.7520e-01, -2.0178e-01,  1.3928e-01,  4.7424e-02,  3.5083e-01,\n",
      "         1.1206e-01,  1.2561e-01,  6.5430e-02,  2.9810e-01,  5.5634e-02,\n",
      "         5.7770e-02,  6.1981e-02,  1.2415e-01,  3.1763e-01,  6.5552e-02,\n",
      "        -1.5674e-01,  3.1616e-02,  2.4451e-01, -1.9189e-01, -5.8398e-01,\n",
      "         3.2166e-02,  2.9565e-01,  1.5308e-01,  3.7781e-02, -2.6831e-01,\n",
      "         4.3018e-01,  1.5540e-01, -2.6709e-01, -4.7705e-01,  9.5459e-02,\n",
      "        -1.9250e-01,  6.8542e-02, -1.0605e+00,  1.0156e-01,  1.0394e-01,\n",
      "         4.2419e-02,  2.1948e-01,  1.8054e-01,  5.1953e-01,  3.7598e-01,\n",
      "        -2.7344e-01,  4.6875e-01,  1.4075e-01,  1.6663e-01, -1.5613e-01,\n",
      "        -4.9530e-02,  1.6516e-01,  1.8738e-01,  1.4258e-01,  1.3403e-01,\n",
      "         1.1255e-01, -1.2201e-01, -1.9861e-01,  3.3447e-01,  1.5808e-01,\n",
      "         1.4160e-01,  1.5125e-01, -1.6724e-01, -8.2947e-02, -2.5342e-01,\n",
      "         2.4048e-02, -2.1216e-01, -4.4403e-02,  4.2297e-02, -1.3660e-01,\n",
      "        -1.9958e-01,  7.8979e-02, -1.9409e-01,  2.3169e-01,  1.4233e-01,\n",
      "         1.9714e-01,  7.6294e-02, -1.2732e-01,  5.0508e+00,  2.7393e-01,\n",
      "        -1.4172e-01, -2.7637e-01,  2.2559e-01,  3.3545e-01, -8.4839e-02,\n",
      "        -4.4824e-01,  2.3523e-01, -9.9121e-02,  9.8938e-02, -2.0374e-01,\n",
      "        -2.0032e-01,  1.6345e-01,  4.9500e-02,  5.9540e-02, -1.7505e-01,\n",
      "         3.9966e-01, -3.9087e-01,  2.1255e-02,  1.3538e-01, -6.7529e-01,\n",
      "        -1.4465e-01,  6.0205e-01,  1.1255e-01, -2.3364e-01,  3.2178e-01,\n",
      "        -3.4424e-01, -3.2623e-02,  2.8003e-01, -3.1143e-02, -3.0908e-01,\n",
      "         1.5076e-01,  1.9739e-01,  3.2837e-01, -1.4282e-01, -4.1284e-01,\n",
      "        -4.1107e-02, -1.0718e-01, -1.8542e-01,  4.1473e-02,  1.5234e-01,\n",
      "         1.8506e-01,  2.1594e-01, -1.6846e-01,  4.2041e-01, -2.3792e-01,\n",
      "        -3.1323e-01,  1.3562e-01, -4.7729e-02, -1.0968e-01, -4.3335e-01,\n",
      "         1.6577e-01,  1.8262e-01, -1.2939e-01,  1.2659e-01,  3.3203e-01,\n",
      "        -9.6283e-03, -2.0523e-02,  5.2002e-01,  4.2603e-02,  5.1318e-01,\n",
      "        -3.4447e-03, -1.1481e-01,  1.9653e-01,  8.0750e-02, -8.6426e-02,\n",
      "         2.1497e-01, -4.0016e-03, -1.1462e-01,  3.7476e-02, -2.1741e-01,\n",
      "         1.5450e-02,  6.0089e-02, -3.7622e-01,  2.7930e-01,  1.3191e-02,\n",
      "         4.6704e-01,  1.1612e-02, -5.3223e-02,  5.2637e-01, -1.5808e-01,\n",
      "         1.4490e-01, -7.3059e-02,  5.5511e-02,  7.3389e-01,  3.4326e-01,\n",
      "        -3.8989e-01,  1.4160e-01, -3.3512e-03,  2.8784e-01, -6.1182e-01,\n",
      "        -2.3880e-02,  1.3501e-01, -4.4212e-03, -1.2524e-01, -3.5205e-01,\n",
      "         1.4197e-01, -2.4796e-03,  3.7384e-02, -1.6577e-01,  3.3179e-01,\n",
      "        -5.4810e-02,  7.2571e-02, -2.1912e-02, -5.1123e-01,  1.2805e-01,\n",
      "         4.4336e-01,  1.1121e-01,  1.8542e-01,  5.8990e-02, -2.1948e-01,\n",
      "         2.4002e-02,  1.1511e-01,  1.0590e-01,  1.9702e-01,  2.1500e-02,\n",
      "         1.2091e-01, -6.3477e-02,  3.5791e-01, -4.0430e-01,  2.5330e-02,\n",
      "        -7.1045e-02,  7.6782e-02, -2.2675e-02, -2.5171e-01, -1.6251e-02,\n",
      "         2.8491e-01,  1.8994e-01, -3.2104e-01, -4.6484e-01, -3.1494e-02,\n",
      "         2.5464e-01,  5.5664e-02,  2.7985e-02,  2.7026e-01, -4.5776e-01,\n",
      "        -1.5283e-01,  1.4172e-01,  6.6345e-02, -2.5101e-02,  6.3286e-03,\n",
      "         3.0957e-01,  4.8950e-01, -3.4106e-01, -3.1830e-02,  1.5249e-03,\n",
      "        -2.2742e-01,  8.2886e-02, -9.6191e-02, -1.0657e-01,  8.1177e-02,\n",
      "        -9.8328e-02, -3.4546e-01,  3.4595e-01,  4.0771e-01, -9.6802e-02,\n",
      "        -5.6343e-03,  1.3647e-01, -1.5747e-01,  9.9731e-02, -1.6003e-03,\n",
      "         1.4763e-02, -5.0439e-01,  3.1030e-01,  4.0649e-01,  7.0557e-02,\n",
      "        -8.2642e-02, -3.7085e-01,  1.2866e-01, -1.4990e-01, -3.7207e-01,\n",
      "         3.4155e-01,  1.5088e-01,  9.7534e-02, -7.3425e-02, -1.9470e-01,\n",
      "         1.3708e-01,  2.0374e-01,  5.0508e+00, -8.1116e-02, -6.2805e-02,\n",
      "         4.0955e-02,  2.3792e-01, -4.0771e-01,  4.8332e-03, -1.5759e-01,\n",
      "         2.1252e-01,  3.8965e-01,  1.5173e-01,  3.8892e-01, -4.1809e-03,\n",
      "         2.9102e-01, -5.9265e-02, -2.9272e-01,  6.5137e-01, -1.1592e+00,\n",
      "         2.3352e-01,  3.9526e-01,  6.1249e-02, -6.2793e-01,  3.7427e-01,\n",
      "         2.0279e-02, -1.2207e-02,  2.5391e-01, -1.3635e-01, -9.3506e-02,\n",
      "        -2.5098e-01, -4.5361e-01, -1.2695e-01, -3.1909e-01, -4.6289e-01,\n",
      "        -3.5425e-01,  4.3274e-02, -2.1484e-01, -1.3318e-01,  6.8604e-02,\n",
      "         6.6223e-02,  1.5796e-01, -2.3303e-01,  3.4448e-01,  3.6816e-01,\n",
      "         3.3521e-01, -2.0740e-01, -3.2983e-01, -6.0254e-01,  8.8928e-02,\n",
      "         2.5122e-01,  2.1741e-01,  9.4666e-02,  2.8296e-01, -1.5723e-01,\n",
      "        -6.9458e-02, -1.5686e-01, -2.1814e-01, -1.3269e-01,  1.9800e-01,\n",
      "        -4.6631e-02, -1.0614e-01,  1.7603e-01,  1.7163e-01, -5.5466e-03,\n",
      "         1.9861e-01,  2.6392e-01, -3.8013e-01, -6.2744e-01,  3.0249e-01,\n",
      "        -4.0552e-01, -9.1699e-01, -1.0858e-01,  2.2751e-02,  1.3025e-01,\n",
      "         3.4863e-01, -4.4800e-01,  1.9714e-01, -2.5146e-01,  2.3926e-01,\n",
      "        -2.2888e-01, -5.4352e-02,  9.4666e-02,  3.1763e-01, -7.2021e-02,\n",
      "        -4.4769e-02,  3.5962e-01, -4.2084e-02, -1.5576e-01,  2.4829e-01,\n",
      "         6.7322e-02,  8.6304e-02, -2.5073e-01, -5.2795e-02,  2.9590e-01,\n",
      "        -2.1423e-01,  4.3823e-02, -1.4168e-02,  3.0176e-01,  4.8309e-02,\n",
      "         9.8022e-02, -9.9121e-02, -6.4636e-02, -3.3325e-01, -4.7516e-02,\n",
      "        -1.0809e-01,  9.1187e-02, -1.5955e-01,  1.4514e-01, -6.4502e-01,\n",
      "         2.5436e-02,  4.3732e-02, -1.8591e-01, -7.3389e-01, -1.0760e-01,\n",
      "        -1.0376e-01,  1.1108e-01, -4.7398e-04, -2.6587e-01, -3.5083e-01,\n",
      "        -8.0200e-02,  2.1448e-01,  7.2815e-02, -1.6260e-01, -3.3228e-01,\n",
      "        -3.7384e-02, -2.6489e-02, -3.4912e-01,  3.1445e-01,  5.8411e-02,\n",
      "         7.9193e-03,  9.1980e-02, -5.7587e-02,  1.3525e-01,  3.1689e-01,\n",
      "         4.6484e-01, -5.4639e-01, -1.9702e-01,  2.6810e-02,  1.9128e-01,\n",
      "         2.3755e-01, -2.5562e-01,  7.7393e-02, -1.1932e-01,  2.0340e-02,\n",
      "        -8.5205e-02,  7.5317e-02, -4.4385e-01, -4.8523e-02, -1.3599e-01,\n",
      "         3.6230e-01, -4.1870e-01,  1.6187e-01, -2.9907e-01, -8.5144e-03,\n",
      "        -5.3741e-02,  1.2070e-02, -3.1885e-01, -4.3945e-01, -2.3987e-01,\n",
      "        -2.5977e-01, -1.5063e-01, -1.5088e-01,  1.1810e-01,  3.4277e-01,\n",
      "        -1.0735e-02,  2.5589e-02, -1.1084e-01, -2.4719e-01, -4.5593e-02,\n",
      "        -1.3513e-01, -7.6953e-01,  1.5308e-01, -6.8726e-02, -9.4849e-02,\n",
      "        -2.4414e-01,  1.2622e-01,  1.5979e-01, -3.7427e-01,  6.7322e-02,\n",
      "         7.9422e-03,  1.2177e-01, -1.2512e-01,  7.7734e-01,  2.8351e-02,\n",
      "        -7.0508e-01,  3.3569e-01,  4.2310e-01, -1.7529e-01,  2.1545e-02,\n",
      "         3.0859e-01, -1.2079e-01, -2.3926e-01,  2.0691e-01,  6.3672e-01,\n",
      "        -5.5786e-02, -3.5675e-02,  4.0820e-01,  3.3618e-01, -2.7173e-01,\n",
      "        -1.9604e-01, -1.4050e-01], dtype=torch.float16)\n",
      "第一个图像嵌入: tensor([-3.0518e-02,  4.7424e-02, -1.9727e-01,  1.3672e-01,  1.8799e-01,\n",
      "        -1.1035e+00, -3.4937e-01, -4.7266e-01,  2.8223e-01,  7.3242e-04,\n",
      "        -1.2817e-01, -3.9038e-01,  5.7520e-01, -2.2400e-02,  6.6943e-01,\n",
      "         1.8262e-01,  1.6406e+00, -4.8364e-01,  5.6671e-02, -2.3962e-01,\n",
      "         1.0406e-01, -2.6416e-01, -3.1958e-01, -3.4619e-01, -2.3010e-01,\n",
      "        -4.9042e-02,  5.8044e-02, -1.6125e-01,  4.4580e-01,  2.3914e-01,\n",
      "         3.3521e-01,  1.2299e-01,  2.5415e-01, -1.1475e-01, -8.8281e-01,\n",
      "         9.9792e-02, -7.9712e-02,  6.1719e-01,  1.5381e-01,  1.6289e+00,\n",
      "        -2.8915e-02,  1.1223e-02, -1.6895e-01, -1.6125e-01, -3.7872e-02,\n",
      "         1.2168e+00,  2.3279e-01,  1.9556e-01,  3.1201e-01,  4.0698e-01,\n",
      "        -1.0175e-01, -4.5654e-02, -9.1858e-02, -3.1274e-01,  3.2568e-01,\n",
      "        -4.7168e-01,  5.0995e-02,  2.0984e-01,  6.8115e-01, -2.8418e-01,\n",
      "         9.8450e-02,  4.4531e-01,  1.8140e-01,  5.1807e-01, -1.7798e-01,\n",
      "         1.7480e-01, -4.2969e-01, -5.1074e-01,  1.5869e-01, -4.1528e-01,\n",
      "         3.2520e-01, -2.2986e-01,  9.5886e-02, -1.7981e-01,  2.3694e-01,\n",
      "        -3.6377e-01, -1.7288e-02, -3.4302e-02, -4.6411e-01,  1.5015e-02,\n",
      "         2.4658e-01,  1.2253e-02,  2.7393e-01, -4.3530e-01,  3.5669e-01,\n",
      "         3.6255e-01,  6.2500e-01,  2.1570e-01, -2.9517e-01,  2.2156e-02,\n",
      "         1.0052e-01, -4.1455e-01, -4.8398e+00, -2.5122e-01, -4.0698e-01,\n",
      "         8.6182e-02,  9.2896e-02, -1.7383e-01, -5.0391e-01,  2.5806e-01,\n",
      "         2.9810e-01, -9.2834e-02, -3.9893e-01, -1.2476e-01,  6.1768e-02,\n",
      "        -1.2866e-01,  1.6572e+00, -3.5864e-01, -1.8848e-01,  5.8350e-01,\n",
      "        -3.3228e-01, -4.5776e-01,  4.2419e-02, -1.1053e-01, -2.6147e-01,\n",
      "         8.1726e-02, -1.3110e-01, -4.5020e-01, -1.1743e-01, -2.1680e-01,\n",
      "        -1.5796e-01, -1.0480e-01,  4.3945e-01,  2.1472e-01, -4.6463e-03,\n",
      "        -2.0142e-02,  3.2715e-01,  1.9128e-01,  3.4088e-02,  5.5176e-01,\n",
      "        -5.9242e-03,  5.1221e-01,  1.3696e-01,  8.1787e-01,  1.5735e-01,\n",
      "         3.9136e-01,  2.9272e-01, -6.3574e-01, -5.7343e-02, -1.4417e-01,\n",
      "         2.0728e-01, -2.2339e-01, -1.8616e-02,  2.5848e-02,  2.2375e-01,\n",
      "         7.3633e-01,  1.6174e-01,  7.6367e-01,  1.2274e-01,  7.7087e-02,\n",
      "        -1.5796e-01,  4.2053e-02,  6.4990e-01, -1.0034e-01,  3.3008e-01,\n",
      "        -3.3252e-01, -4.3311e-01,  1.9434e-01, -5.7678e-03, -2.6538e-01,\n",
      "        -2.0642e-01, -3.9209e-01,  2.3401e-01,  1.8127e-01, -3.9062e-01,\n",
      "         2.7124e-01,  1.5564e-01, -4.1504e-02,  7.1526e-03,  1.3562e-01,\n",
      "        -1.9913e-02,  1.3647e-01,  2.3413e-01,  1.4856e-01, -1.7310e-01,\n",
      "         4.8071e-01,  2.4341e-01,  3.1763e-01,  5.5811e-01,  5.6519e-02,\n",
      "        -2.1509e-01, -2.2009e-01,  3.5864e-01, -1.4124e-01, -1.1975e-01,\n",
      "        -2.0020e-02, -1.8921e-01, -7.2975e-03, -8.7830e-02,  2.8345e-01,\n",
      "        -3.5553e-03, -6.4880e-02,  3.4619e-01,  1.6052e-01, -2.5269e-01,\n",
      "         4.5850e-01, -4.9896e-02,  4.8193e-01,  3.3716e-01,  7.4890e-02,\n",
      "        -7.7515e-02, -3.0005e-01, -5.3516e-01, -1.1945e-01,  4.2334e-01,\n",
      "        -8.7109e-01,  1.1902e-01, -8.5205e-02,  7.1631e-01, -5.5481e-02,\n",
      "         4.6167e-01,  2.2620e-01,  2.4023e-01, -7.6477e-02, -7.7332e-02,\n",
      "         2.9175e-02, -7.4463e-02,  3.2373e-01,  1.1816e+00, -1.1023e-01,\n",
      "        -3.3594e-01, -4.7412e-01, -6.8457e-01, -2.0837e-01, -3.7646e-01,\n",
      "         1.3464e-01, -6.7383e-02, -6.2012e-01, -1.2006e-01, -1.7065e-01,\n",
      "        -8.3740e-02, -4.5581e-01,  1.2830e-01, -1.4062e-01,  2.7539e-01,\n",
      "        -8.7012e-01, -2.7686e-01,  2.8101e-01, -7.1240e-01,  1.4978e-01,\n",
      "        -2.5122e-01,  3.8281e-01, -5.6671e-02,  7.7057e-03, -4.4769e-02,\n",
      "         2.4084e-01,  7.8674e-02,  2.1362e-01, -3.4515e-02,  1.5564e-01,\n",
      "         2.5000e-01,  4.5264e-01,  3.9966e-01,  2.3425e-01,  5.4297e-01,\n",
      "         2.8052e-01, -4.0088e-01, -4.9164e-02,  1.0352e-01,  4.5715e-02,\n",
      "         1.4844e-01, -1.4868e-01,  7.9712e-02, -3.8745e-01, -1.3074e-01,\n",
      "         4.1797e-01, -5.0385e-02, -1.4420e-02,  2.0508e-01, -2.9004e-01,\n",
      "         9.6802e-02,  2.0251e-01,  8.8623e-02, -1.2708e-01,  8.5632e-02,\n",
      "         2.1790e-01,  1.8311e-01,  2.1887e-01,  5.9180e-01, -5.5518e-01,\n",
      "         2.8003e-01, -4.3671e-02, -1.1414e-01, -2.7637e-01,  1.5771e-01,\n",
      "         5.5127e-01, -2.3401e-01, -7.1350e-02, -2.2864e-01,  1.3464e-01,\n",
      "        -6.1737e-02, -1.1948e-02, -6.3660e-02,  4.3042e-01,  3.3356e-02,\n",
      "        -1.5344e-01, -1.3262e+00,  3.2074e-02,  5.5225e-01,  8.3130e-02,\n",
      "        -5.9277e-01, -4.9927e-01, -3.9429e-01, -8.7280e-02, -1.0120e-01,\n",
      "         5.1416e-01, -1.1353e-01, -1.7615e-01, -3.5425e-01,  2.3087e-02,\n",
      "         6.8542e-02, -5.7800e-02,  8.1445e-01, -2.4994e-02, -2.4487e-01,\n",
      "         2.5366e-01, -7.8430e-02,  1.7200e-01, -6.1226e-03, -6.9287e-01,\n",
      "         3.0298e-01,  1.4429e-01, -3.4106e-01,  3.6792e-01, -1.2927e-01,\n",
      "         1.5198e-01, -8.4167e-02,  4.6973e-01,  9.1992e-01, -4.8340e-01,\n",
      "        -2.2510e-01, -5.5469e-01,  3.2544e-01, -5.7373e-02, -3.1738e-01,\n",
      "         9.7168e-02, -1.4661e-01, -1.0400e-01,  9.4604e-03, -9.2468e-02,\n",
      "        -1.6455e-01, -2.0142e-02, -2.1118e-02, -1.4722e-01, -3.6255e-01,\n",
      "        -1.9568e-01, -1.8762e-01,  1.0590e-01, -2.2119e-01, -2.7246e-01,\n",
      "         1.1982e+00,  8.7830e-02,  2.2815e-01, -2.5513e-01,  3.6963e-01,\n",
      "        -1.0498e+00, -4.2969e-01,  5.7080e-01, -1.8677e-01,  3.2129e-01,\n",
      "         8.6426e-02, -7.1594e-02, -3.1323e-01, -7.1411e-02, -1.5361e+00,\n",
      "         2.0422e-01,  1.1890e-01, -1.3105e+00,  5.4932e-01, -1.5454e-01,\n",
      "        -4.0747e-01, -2.4146e-01,  4.3774e-01, -4.4409e-01,  1.8005e-01,\n",
      "        -7.2144e-02,  8.2129e-01,  3.2013e-02,  2.7588e-01, -3.1055e-01,\n",
      "        -4.1309e-01,  2.3865e-01, -3.3618e-01, -3.7689e-02,  6.8542e-02,\n",
      "        -7.5417e-03,  2.4390e-01, -4.7974e-02, -1.9299e-01,  1.8379e+00,\n",
      "        -3.6841e-01, -2.6831e-01, -1.0414e-02,  3.4082e-01,  5.7520e-01,\n",
      "         1.6602e-01, -1.9165e-02, -3.5229e-01,  1.7004e-01, -6.4307e-01,\n",
      "        -2.8345e-01, -3.3618e-01,  5.9375e-01, -1.7334e-02,  3.2446e-01,\n",
      "         6.2866e-02,  7.7637e-02, -2.9419e-01,  3.0396e-01,  2.3315e-02,\n",
      "         2.3560e-01,  2.3828e-01,  1.4075e-01,  1.2512e-01, -8.6914e-02,\n",
      "        -2.1094e-01, -1.3745e-01,  2.6709e-01, -6.9482e-01, -6.1377e-01,\n",
      "         5.3619e-02, -1.9421e-01, -2.2107e-01, -4.0820e-01, -2.1313e-01,\n",
      "        -1.8469e-01,  2.3584e-01, -2.5513e-02, -8.7109e-01, -2.0239e-01,\n",
      "        -1.0034e-01,  6.8604e-01,  4.2163e-01, -2.8247e-01,  1.2817e-01,\n",
      "         2.5269e-01,  3.8525e-01,  1.4722e-01, -1.3428e-01,  1.0718e-01,\n",
      "        -2.9712e-01,  9.8572e-02, -3.0835e-01,  5.8167e-02,  1.4136e-01,\n",
      "         4.0550e-03, -8.5449e-02, -4.0942e-01,  2.0483e-01, -5.0293e-01,\n",
      "         2.5195e-01,  1.1572e-01, -3.3398e-01, -3.3569e-02, -1.7627e-01,\n",
      "        -1.3342e-01,  3.5962e-01,  8.0139e-02, -1.7151e-01, -4.1553e-01,\n",
      "         3.7573e-01,  4.1553e-01,  4.2676e-01, -2.2363e-01, -2.6636e-01,\n",
      "         3.0444e-01, -2.5317e-01,  2.7905e-01,  1.5479e-01,  1.3257e-01,\n",
      "         6.3110e-02,  1.4526e-01,  6.5247e-02, -3.7109e-02,  1.1359e-01,\n",
      "         2.3914e-01,  4.1528e-01,  2.0642e-01,  1.0449e-01,  1.3464e-01,\n",
      "         2.7197e-01, -3.5791e-01,  2.5854e-01, -7.2083e-02, -1.4148e-01,\n",
      "         7.9834e-02, -3.8306e-01,  2.7075e-01, -1.3843e-01,  4.3018e-01,\n",
      "        -1.9470e-01, -3.3618e-01,  4.4116e-01, -1.1877e-01, -5.0537e-02,\n",
      "        -4.4946e-01, -1.1487e-01,  2.5684e-01,  1.0339e-01,  8.3618e-02,\n",
      "         2.0862e-01,  1.2903e-01,  6.8604e-02,  8.7891e-02, -2.9861e-02,\n",
      "        -5.8301e-01,  3.6597e-01, -9.4116e-02, -3.3496e-01,  6.8799e-01,\n",
      "         2.9007e-02,  1.0786e-03], dtype=torch.float16)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "text_emb_sample = torch.load('./clip_embeddings/text_embeddings.pt')\n",
    "image_emb_sample = torch.load('./clip_embeddings/image_embeddings.pt')\n",
    "\n",
    "print(f\"文本嵌入样本形状: {text_emb_sample.shape}\")  # 应为 [N, 512]\n",
    "print(f\"图像嵌入样本形状: {image_emb_sample.shape}\")  # 应为 [N, 512]\n",
    "\n",
    "# 打印第一个样本的嵌入\n",
    "print(\"第一个文本嵌入:\", text_emb_sample[0])\n",
    "print(\"第一个图像嵌入:\", image_emb_sample[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 下面这个是新的嵌入生成，是我们为了训练和测试过程都用图像嵌入去运行代码（这样效果一定很好，但是违规哟，测试用文本才可以）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cuda\n",
      "Attributes DataFrame columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n",
      "Attributes DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Sideburns  Smiling  \\\n",
      "0    -1     -1        -1        -1          -1  ...         -1        1   \n",
      "1    -1     -1        -1         1          -1  ...         -1        1   \n",
      "2    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "3    -1     -1        -1        -1          -1  ...         -1       -1   \n",
      "4    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "\n",
      "   Straight_Hair  Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0              1         -1                 1           -1                 1   \n",
      "1             -1         -1                -1           -1                -1   \n",
      "2             -1          1                -1           -1                -1   \n",
      "3              1         -1                 1           -1                 1   \n",
      "4             -1         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  \n",
      "0                -1               -1      1  \n",
      "1                -1               -1      1  \n",
      "2                -1               -1      1  \n",
      "3                 1               -1      1  \n",
      "4                -1               -1      1  \n",
      "\n",
      "[5 rows x 41 columns]\n",
      "'image_id' column found in attributes DataFrame.\n",
      "Partition DataFrame columns: ['image_id', 'partition']\n",
      "Partition DataFrame head:\n",
      "      image_id  partition\n",
      "0  000001.jpg          0\n",
      "1  000002.jpg          0\n",
      "2  000003.jpg          0\n",
      "3  000004.jpg          0\n",
      "4  000005.jpg          0\n",
      "Bounding Box DataFrame columns: ['image_id', 'x_1', 'y_1', 'width', 'height']\n",
      "Bounding Box DataFrame head:\n",
      "      image_id  x_1  y_1  width  height\n",
      "0  000001.jpg   95   71    226     313\n",
      "1  000002.jpg   72   94    221     306\n",
      "2  000003.jpg  216   59     91     126\n",
      "3  000004.jpg  622  257    564     781\n",
      "4  000005.jpg  236  109    120     166\n",
      "After merging with partition, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition']\n",
      "Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Smiling  Straight_Hair  \\\n",
      "0    -1     -1        -1        -1          -1  ...        1              1   \n",
      "1    -1     -1        -1         1          -1  ...        1             -1   \n",
      "2    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "3    -1     -1        -1        -1          -1  ...       -1              1   \n",
      "4    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "\n",
      "   Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0         -1                 1           -1                 1   \n",
      "1         -1                -1           -1                -1   \n",
      "2          1                -1           -1                -1   \n",
      "3         -1                 1           -1                 1   \n",
      "4         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  partition  \n",
      "0                -1               -1      1          0  \n",
      "1                -1               -1      1          0  \n",
      "2                -1               -1      1          0  \n",
      "3                 1               -1      1          0  \n",
      "4                -1               -1      1          0  \n",
      "\n",
      "[5 rows x 42 columns]\n",
      "After merging with bounding box, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition', 'x_1', 'y_1', 'width', 'height']\n",
      "Final Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Wearing_Hat  \\\n",
      "0    -1     -1        -1        -1          -1  ...           -1   \n",
      "1    -1     -1        -1         1          -1  ...           -1   \n",
      "2    -1     -1         1        -1          -1  ...           -1   \n",
      "3    -1     -1        -1        -1          -1  ...           -1   \n",
      "4    -1     -1         1        -1          -1  ...           -1   \n",
      "\n",
      "   Wearing_Lipstick  Wearing_Necklace  Wearing_Necktie  Young  partition  x_1  \\\n",
      "0                 1                -1               -1      1          0   95   \n",
      "1                -1                -1               -1      1          0   72   \n",
      "2                -1                -1               -1      1          0  216   \n",
      "3                 1                 1               -1      1          0  622   \n",
      "4                 1                -1               -1      1          0  236   \n",
      "\n",
      "   y_1  width  height  \n",
      "0   71    226     313  \n",
      "1   94    221     306  \n",
      "2   59     91     126  \n",
      "3  257    564     781  \n",
      "4  109    120     166  \n",
      "\n",
      "[5 rows x 46 columns]\n",
      "Total samples before filtering: 202599\n",
      "Total samples after filtering partition 0: 162770\n",
      "Attribute names: ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings for partition 0:   0%|          | 0/318 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index: 0, Text Prompt: A person with Arched_Eyebrows, Attractive, Brown_Hair, Heavy_Makeup, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Pointy_Nose, Smiling, Straight_Hair.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 1, Text Prompt: A person with Bags_Under_Eyes, Big_Nose, Brown_Hair, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Smiling, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 2, Text Prompt: A person with Big_Lips, Blurry, Male, Narrow_Eyes, No_Beard, Pointy_Nose, Wavy_Hair, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 3, Text Prompt: A person with Attractive, No_Beard, Pointy_Nose, Straight_Hair, Wearing_Earrings, Wearing_Lipstick, Wearing_Necklace, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 4, Text Prompt: A person with Arched_Eyebrows, Attractive, Big_Lips, Heavy_Makeup, Narrow_Eyes, No_Beard, Pointy_Nose, Wearing_Lipstick, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings for partition 0: 100%|██████████| 318/318 [02:17<00:00,  2.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "嵌入已保存为 'text_embeddings_partition_0.pt' 和 'image_embeddings_partition_0.pt'\n",
      "Attributes DataFrame columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n",
      "Attributes DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Sideburns  Smiling  \\\n",
      "0    -1     -1        -1        -1          -1  ...         -1        1   \n",
      "1    -1     -1        -1         1          -1  ...         -1        1   \n",
      "2    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "3    -1     -1        -1        -1          -1  ...         -1       -1   \n",
      "4    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "\n",
      "   Straight_Hair  Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0              1         -1                 1           -1                 1   \n",
      "1             -1         -1                -1           -1                -1   \n",
      "2             -1          1                -1           -1                -1   \n",
      "3              1         -1                 1           -1                 1   \n",
      "4             -1         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  \n",
      "0                -1               -1      1  \n",
      "1                -1               -1      1  \n",
      "2                -1               -1      1  \n",
      "3                 1               -1      1  \n",
      "4                -1               -1      1  \n",
      "\n",
      "[5 rows x 41 columns]\n",
      "'image_id' column found in attributes DataFrame.\n",
      "Partition DataFrame columns: ['image_id', 'partition']\n",
      "Partition DataFrame head:\n",
      "      image_id  partition\n",
      "0  000001.jpg          0\n",
      "1  000002.jpg          0\n",
      "2  000003.jpg          0\n",
      "3  000004.jpg          0\n",
      "4  000005.jpg          0\n",
      "Bounding Box DataFrame columns: ['image_id', 'x_1', 'y_1', 'width', 'height']\n",
      "Bounding Box DataFrame head:\n",
      "      image_id  x_1  y_1  width  height\n",
      "0  000001.jpg   95   71    226     313\n",
      "1  000002.jpg   72   94    221     306\n",
      "2  000003.jpg  216   59     91     126\n",
      "3  000004.jpg  622  257    564     781\n",
      "4  000005.jpg  236  109    120     166\n",
      "After merging with partition, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition']\n",
      "Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Smiling  Straight_Hair  \\\n",
      "0    -1     -1        -1        -1          -1  ...        1              1   \n",
      "1    -1     -1        -1         1          -1  ...        1             -1   \n",
      "2    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "3    -1     -1        -1        -1          -1  ...       -1              1   \n",
      "4    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "\n",
      "   Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0         -1                 1           -1                 1   \n",
      "1         -1                -1           -1                -1   \n",
      "2          1                -1           -1                -1   \n",
      "3         -1                 1           -1                 1   \n",
      "4         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  partition  \n",
      "0                -1               -1      1          0  \n",
      "1                -1               -1      1          0  \n",
      "2                -1               -1      1          0  \n",
      "3                 1               -1      1          0  \n",
      "4                -1               -1      1          0  \n",
      "\n",
      "[5 rows x 42 columns]\n",
      "After merging with bounding box, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition', 'x_1', 'y_1', 'width', 'height']\n",
      "Final Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Wearing_Hat  \\\n",
      "0    -1     -1        -1        -1          -1  ...           -1   \n",
      "1    -1     -1        -1         1          -1  ...           -1   \n",
      "2    -1     -1         1        -1          -1  ...           -1   \n",
      "3    -1     -1        -1        -1          -1  ...           -1   \n",
      "4    -1     -1         1        -1          -1  ...           -1   \n",
      "\n",
      "   Wearing_Lipstick  Wearing_Necklace  Wearing_Necktie  Young  partition  x_1  \\\n",
      "0                 1                -1               -1      1          0   95   \n",
      "1                -1                -1               -1      1          0   72   \n",
      "2                -1                -1               -1      1          0  216   \n",
      "3                 1                 1               -1      1          0  622   \n",
      "4                 1                -1               -1      1          0  236   \n",
      "\n",
      "   y_1  width  height  \n",
      "0   71    226     313  \n",
      "1   94    221     306  \n",
      "2   59     91     126  \n",
      "3  257    564     781  \n",
      "4  109    120     166  \n",
      "\n",
      "[5 rows x 46 columns]\n",
      "Total samples before filtering: 202599\n",
      "Total samples after filtering partition 1: 19867\n",
      "Attribute names: ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings for partition 1:   0%|          | 0/39 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index: 0, Text Prompt: A person with Attractive, Bangs, Big_Lips, Heavy_Makeup, High_Cheekbones, No_Beard, Oval_Face, Pointy_Nose, Smiling, Wavy_Hair.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 1, Text Prompt: A person with Big_Nose, Black_Hair, Bushy_Eyebrows, Goatee, Male, Mouth_Slightly_Open, Mustache, Pointy_Nose, Sideburns, Smiling.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 2, Text Prompt: A person with Arched_Eyebrows, Attractive, Big_Lips, Brown_Hair, Heavy_Makeup, No_Beard, Pointy_Nose, Wearing_Lipstick, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 3, Text Prompt: A person with 5_o_Clock_Shadow, Attractive, Brown_Hair, Male, Mouth_Slightly_Open, Oval_Face, Pointy_Nose, Receding_Hairline, Smiling, Straight_Hair.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 4, Text Prompt: A person with Black_Hair, No_Beard, Pointy_Nose, Wearing_Lipstick, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings for partition 1: 100%|██████████| 39/39 [00:20<00:00,  1.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "嵌入已保存为 'text_embeddings_partition_1.pt' 和 'image_embeddings_partition_1.pt'\n",
      "Attributes DataFrame columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n",
      "Attributes DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Sideburns  Smiling  \\\n",
      "0    -1     -1        -1        -1          -1  ...         -1        1   \n",
      "1    -1     -1        -1         1          -1  ...         -1        1   \n",
      "2    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "3    -1     -1        -1        -1          -1  ...         -1       -1   \n",
      "4    -1     -1         1        -1          -1  ...         -1       -1   \n",
      "\n",
      "   Straight_Hair  Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0              1         -1                 1           -1                 1   \n",
      "1             -1         -1                -1           -1                -1   \n",
      "2             -1          1                -1           -1                -1   \n",
      "3              1         -1                 1           -1                 1   \n",
      "4             -1         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  \n",
      "0                -1               -1      1  \n",
      "1                -1               -1      1  \n",
      "2                -1               -1      1  \n",
      "3                 1               -1      1  \n",
      "4                -1               -1      1  \n",
      "\n",
      "[5 rows x 41 columns]\n",
      "'image_id' column found in attributes DataFrame.\n",
      "Partition DataFrame columns: ['image_id', 'partition']\n",
      "Partition DataFrame head:\n",
      "      image_id  partition\n",
      "0  000001.jpg          0\n",
      "1  000002.jpg          0\n",
      "2  000003.jpg          0\n",
      "3  000004.jpg          0\n",
      "4  000005.jpg          0\n",
      "Bounding Box DataFrame columns: ['image_id', 'x_1', 'y_1', 'width', 'height']\n",
      "Bounding Box DataFrame head:\n",
      "      image_id  x_1  y_1  width  height\n",
      "0  000001.jpg   95   71    226     313\n",
      "1  000002.jpg   72   94    221     306\n",
      "2  000003.jpg  216   59     91     126\n",
      "3  000004.jpg  622  257    564     781\n",
      "4  000005.jpg  236  109    120     166\n",
      "After merging with partition, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition']\n",
      "Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Smiling  Straight_Hair  \\\n",
      "0    -1     -1        -1        -1          -1  ...        1              1   \n",
      "1    -1     -1        -1         1          -1  ...        1             -1   \n",
      "2    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "3    -1     -1        -1        -1          -1  ...       -1              1   \n",
      "4    -1     -1         1        -1          -1  ...       -1             -1   \n",
      "\n",
      "   Wavy_Hair  Wearing_Earrings  Wearing_Hat  Wearing_Lipstick  \\\n",
      "0         -1                 1           -1                 1   \n",
      "1         -1                -1           -1                -1   \n",
      "2          1                -1           -1                -1   \n",
      "3         -1                 1           -1                 1   \n",
      "4         -1                -1           -1                 1   \n",
      "\n",
      "   Wearing_Necklace  Wearing_Necktie  Young  partition  \n",
      "0                -1               -1      1          0  \n",
      "1                -1               -1      1          0  \n",
      "2                -1               -1      1          0  \n",
      "3                 1               -1      1          0  \n",
      "4                -1               -1      1          0  \n",
      "\n",
      "[5 rows x 42 columns]\n",
      "After merging with bounding box, columns: ['image_id', '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', 'partition', 'x_1', 'y_1', 'width', 'height']\n",
      "Final Merged DataFrame head:\n",
      "      image_id  5_o_Clock_Shadow  Arched_Eyebrows  Attractive  Bags_Under_Eyes  \\\n",
      "0  000001.jpg                -1                1           1               -1   \n",
      "1  000002.jpg                -1               -1          -1                1   \n",
      "2  000003.jpg                -1               -1          -1               -1   \n",
      "3  000004.jpg                -1               -1           1               -1   \n",
      "4  000005.jpg                -1                1           1               -1   \n",
      "\n",
      "   Bald  Bangs  Big_Lips  Big_Nose  Black_Hair  ...  Wearing_Hat  \\\n",
      "0    -1     -1        -1        -1          -1  ...           -1   \n",
      "1    -1     -1        -1         1          -1  ...           -1   \n",
      "2    -1     -1         1        -1          -1  ...           -1   \n",
      "3    -1     -1        -1        -1          -1  ...           -1   \n",
      "4    -1     -1         1        -1          -1  ...           -1   \n",
      "\n",
      "   Wearing_Lipstick  Wearing_Necklace  Wearing_Necktie  Young  partition  x_1  \\\n",
      "0                 1                -1               -1      1          0   95   \n",
      "1                -1                -1               -1      1          0   72   \n",
      "2                -1                -1               -1      1          0  216   \n",
      "3                 1                 1               -1      1          0  622   \n",
      "4                 1                -1               -1      1          0  236   \n",
      "\n",
      "   y_1  width  height  \n",
      "0   71    226     313  \n",
      "1   94    221     306  \n",
      "2   59     91     126  \n",
      "3  257    564     781  \n",
      "4  109    120     166  \n",
      "\n",
      "[5 rows x 46 columns]\n",
      "Total samples before filtering: 202599\n",
      "Total samples after filtering partition 2: 19962\n",
      "Attribute names: ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', 'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings for partition 2:   0%|          | 0/39 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index: 0, Text Prompt: A person with Arched_Eyebrows, Bags_Under_Eyes, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Pointy_Nose, Smiling, Wavy_Hair, Wearing_Earrings, Wearing_Hat.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 1, Text Prompt: A person with Big_Nose, Black_Hair, Bushy_Eyebrows, High_Cheekbones, Narrow_Eyes, No_Beard, Wavy_Hair, Wearing_Lipstick, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 2, Text Prompt: A person with Attractive, Brown_Hair, Heavy_Makeup, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Oval_Face, Smiling, Wavy_Hair, Wearing_Lipstick.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 3, Text Prompt: A person with Attractive, Heavy_Makeup, High_Cheekbones, Mouth_Slightly_Open, No_Beard, Oval_Face, Pointy_Nose, Rosy_Cheeks, Smiling, Wearing_Earrings.\n",
      "Image shape: torch.Size([3, 224, 224])\n",
      "Index: 4, Text Prompt: A person with Black_Hair, High_Cheekbones, Mouth_Slightly_Open, Narrow_Eyes, No_Beard, Smiling, Young.\n",
      "Image shape: torch.Size([3, 224, 224])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing Embeddings for partition 2: 100%|██████████| 39/39 [00:23<00:00,  1.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "嵌入已保存为 'text_embeddings_partition_2.pt' 和 'image_embeddings_partition_2.pt'\n"
     ]
    }
   ],
   "source": [
    "# preprocess_clip_embeddings.py\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import clip\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "# 配置参数\n",
    "IMG_DIR = '/root/autodl-tmp/celeba_datasets/img_align_celeba/img_align_celeba'\n",
    "ATTR_PATH = '/root/autodl-tmp/celeba_datasets/list_attr_celeba.txt'\n",
    "BBOX_PATH = '/root/autodl-tmp/celeba_datasets/list_bbox_celeba.txt'\n",
    "PARTITION_PATH = '/root/autodl-tmp/celeba_datasets/list_eval_partition.txt'\n",
    "\n",
    "OUTPUT_DIR = './clip_embeddings'  # 输出目录\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "IMAGE_SIZE = 224  # CLIP ViT-B/32的默认图像尺寸\n",
    "BATCH_SIZE = 512  # 根据GPU显存调整\n",
    "NUM_WORKERS = 4  # 根据CPU核心数调整\n",
    "\n",
    "# 设备设置\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "# 图像预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.48145466, 0.4578275, 0.40821073], \n",
    "                         [0.26862954, 0.26130258, 0.27577711])  # CLIP的归一化参数\n",
    "])\n",
    "\n",
    "# 数据集类\n",
    "class CelebAEmbeddingsDataset(Dataset):\n",
    "    def __init__(self, img_dir, attr_path, bbox_path, partition_path, transform=None, partition=0):\n",
    "        self.img_dir = img_dir\n",
    "        self.transform = transform\n",
    "\n",
    "        # 读取属性文件\n",
    "        try:\n",
    "            attr_df = pd.read_csv(attr_path, sep=',', header=0)\n",
    "        except Exception as e:\n",
    "            print(f\"Error reading {attr_path}: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"Attributes DataFrame columns:\", attr_df.columns.tolist())\n",
    "        print(\"Attributes DataFrame head:\\n\", attr_df.head())\n",
    "\n",
    "        # 确认 'image_id' 列存在\n",
    "        if 'image_id' not in attr_df.columns:\n",
    "            raise ValueError(\"属性文件中缺少 'image_id' 列。请检查文件格式。\")\n",
    "        else:\n",
    "            print(\"'image_id' column found in attributes DataFrame.\")\n",
    "\n",
    "        # 读取分区文件\n",
    "        try:\n",
    "            partition_df = pd.read_csv(partition_path, sep=',', header=0, names=['image_id', 'partition'])\n",
    "        except Exception as e:\n",
    "            print(f\"Error reading {partition_path}: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"Partition DataFrame columns:\", partition_df.columns.tolist())\n",
    "        print(\"Partition DataFrame head:\\n\", partition_df.head())\n",
    "\n",
    "        # 读取边界框文件\n",
    "        try:\n",
    "            bbox_df = pd.read_csv(bbox_path, sep=',', header=0, names=['image_id', 'x_1', 'y_1', 'width', 'height'])\n",
    "        except Exception as e:\n",
    "            print(f\"Error reading {bbox_path}: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"Bounding Box DataFrame columns:\", bbox_df.columns.tolist())\n",
    "        print(\"Bounding Box DataFrame head:\\n\", bbox_df.head())\n",
    "\n",
    "        # 合并属性和分区数据\n",
    "        try:\n",
    "            attr_df = attr_df.merge(partition_df, on='image_id')\n",
    "        except Exception as e:\n",
    "            print(f\"Error merging attributes and partitions: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"After merging with partition, columns:\", attr_df.columns.tolist())\n",
    "        print(\"Merged DataFrame head:\\n\", attr_df.head())\n",
    "\n",
    "        # 合并属性、分区和边界框数据\n",
    "        try:\n",
    "            attr_df = attr_df.merge(bbox_df, on='image_id')\n",
    "        except Exception as e:\n",
    "            print(f\"Error merging attributes, partitions and bounding boxes: {e}\")\n",
    "            raise\n",
    "\n",
    "        print(\"After merging with bounding box, columns:\", attr_df.columns.tolist())\n",
    "        print(\"Final Merged DataFrame head:\\n\", attr_df.head())\n",
    "\n",
    "        # 筛选分区\n",
    "        print(f\"Total samples before filtering: {len(attr_df)}\")\n",
    "        self.attr_df = attr_df[attr_df['partition'] == partition].reset_index(drop=True)\n",
    "        print(f\"Total samples after filtering partition {partition}: {len(self.attr_df)}\")\n",
    "\n",
    "        if len(self.attr_df) == 0:\n",
    "            raise ValueError(f\"训练集中没有样本。请检查数据文件的读取和过滤逻辑。\")\n",
    "\n",
    "        # 获取属性名称\n",
    "        # 属性列是除了 'image_id', 'partition', 'x_1', 'y_1', 'width', 'height' 之外的所有列\n",
    "        self.attr_names = [col for col in attr_df.columns if col not in ['image_id', 'partition', 'x_1', 'y_1', 'width', 'height']]\n",
    "        print(\"Attribute names:\", self.attr_names)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.attr_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # 获取图像路径\n",
    "        img_name = self.attr_df.iloc[idx, self.attr_df.columns.get_loc('image_id')]\n",
    "        img_path = os.path.join(self.img_dir, img_name)\n",
    "        try:\n",
    "            image = Image.open(img_path).convert('RGB')\n",
    "        except Exception as e:\n",
    "            print(f\"Error opening image {img_path}: {e}\")\n",
    "            # 可以选择跳过该样本，或返回一个默认图像\n",
    "            # 这里选择抛出异常\n",
    "            raise\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        # 获取属性标签并生成文本描述\n",
    "        attrs = self.attr_df.iloc[idx][self.attr_names].values\n",
    "        attrs = (attrs + 1) // 2  # 将-1转为0，1保持1\n",
    "        attrs = attrs.astype(int)\n",
    "        active_attrs = [attr for attr, val in zip(self.attr_names, attrs) if val == 1]\n",
    "\n",
    "        # 手动截断属性数量，例如，最多10个属性，以避免超长文本\n",
    "        max_attrs = 10\n",
    "        active_attrs = active_attrs[:max_attrs]\n",
    "\n",
    "        if len(active_attrs) == 0:\n",
    "            text_prompt = \"A person with no notable attributes.\"\n",
    "        else:\n",
    "            text_description = ', '.join(active_attrs)\n",
    "            text_prompt = f\"A person with {text_description}.\"\n",
    "\n",
    "        # 打印前5个样本的文本描述和图像形状\n",
    "        if idx < 5:\n",
    "            print(f\"Index: {idx}, Text Prompt: {text_prompt}\")\n",
    "            print(f\"Image shape: {image.shape}\")\n",
    "\n",
    "        return text_prompt, image\n",
    "\n",
    "# 初始化 CLIP 模型\n",
    "clip_model, preprocess_clip = clip.load(\"ViT-B/32\", device=device)\n",
    "clip_model.eval()\n",
    "\n",
    "# 创建数据集和数据加载器，并计算嵌入\n",
    "def compute_and_save_embeddings(partition, output_dir):\n",
    "    dataset = CelebAEmbeddingsDataset(\n",
    "        img_dir=IMG_DIR,\n",
    "        attr_path=ATTR_PATH,\n",
    "        bbox_path=BBOX_PATH,\n",
    "        partition_path=PARTITION_PATH,\n",
    "        transform=transform,\n",
    "        partition=partition\n",
    "    )\n",
    "    dataloader = DataLoader(\n",
    "        dataset,\n",
    "        batch_size=BATCH_SIZE,\n",
    "        shuffle=False,\n",
    "        num_workers=NUM_WORKERS,\n",
    "        pin_memory=True\n",
    "    )\n",
    "\n",
    "    text_embeddings = []\n",
    "    image_embeddings = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for text_prompts, images in tqdm(dataloader, desc=f\"Computing Embeddings for partition {partition}\"):\n",
    "            # 生成文本嵌入\n",
    "            try:\n",
    "                text_tokens = clip.tokenize(text_prompts, truncate=True).to(device)\n",
    "            except Exception as e:\n",
    "                print(f\"Error tokenizing text: {e}\")\n",
    "                raise\n",
    "\n",
    "            try:\n",
    "                text_emb = clip_model.encode_text(text_tokens).cpu()\n",
    "            except Exception as e:\n",
    "                print(f\"Error encoding text: {e}\")\n",
    "                raise\n",
    "\n",
    "            text_embeddings.append(text_emb)\n",
    "\n",
    "            # 生成图像嵌入\n",
    "            try:\n",
    "                images = images.to(device)\n",
    "                image_emb = clip_model.encode_image(images).cpu()\n",
    "            except Exception as e:\n",
    "                print(f\"Error encoding images: {e}\")\n",
    "                raise\n",
    "\n",
    "            image_embeddings.append(image_emb)\n",
    "\n",
    "    # 合并所有嵌入\n",
    "    if len(text_embeddings) == 0 or len(image_embeddings) == 0:\n",
    "        raise RuntimeError(\"未计算到任何嵌入。请检查数据加载和嵌入计算步骤。\")\n",
    "\n",
    "    text_embeddings = torch.cat(text_embeddings, dim=0)\n",
    "    image_embeddings = torch.cat(image_embeddings, dim=0)\n",
    "\n",
    "    # 保存嵌入到文件\n",
    "    torch.save(text_embeddings, os.path.join(output_dir, f'text_embeddings_partition_{partition}.pt'))\n",
    "    torch.save(image_embeddings, os.path.join(output_dir, f'image_embeddings_partition_{partition}.pt'))\n",
    "    print(f\"嵌入已保存为 'text_embeddings_partition_{partition}.pt' 和 'image_embeddings_partition_{partition}.pt'\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # 分别计算并保存每个分区的嵌入\n",
    "    for partition in [0, 1, 2]:  # 0: train, 1: val, 2: test\n",
    "        compute_and_save_embeddings(partition, OUTPUT_DIR)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorc_test1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "undefined.undefined.undefined"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}