text_to_image_embedder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练 CLIP 文本到图像嵌入的映射网络"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. 环境设置\n",
    "首先，导入必要的库，并配置一些基本参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cuda\n"
     ]
    }
   ],
   "source": [
    "# 导入必要的库\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "# 设置随机种子以确保可重复性\n",
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "random.seed(42)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(42)\n",
    "\n",
    "# 配置参数\n",
    "TEXT_EMBEDDINGS_PATH = './clip_embeddings/text_embeddings.pt'   # 替换为实际路径\n",
    "IMAGE_EMBEDDINGS_PATH = './clip_embeddings/image_embeddings.pt' # 替换为实际路径\n",
    "MODEL_SAVE_PATH = 'text_to_image_embedder.pth'\n",
    "BATCH_SIZE = 128\n",
    "NUM_EPOCHS = 10\n",
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "# 设备设置\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. 定义数据集类\n",
    "创建一个自定义的 Dataset 类，用于加载预先计算的文本和图像嵌入。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义自定义 Dataset 类\n",
    "class CLIPEmbedMappingDatasetPrecomputed(Dataset):\n",
    "    def __init__(self, text_embeddings_path, image_embeddings_path):\n",
    "        \"\"\"\n",
    "        初始化数据集，加载预计算的嵌入\n",
    "        :param text_embeddings_path: 文本嵌入文件路径\n",
    "        :param image_embeddings_path: 图像嵌入文件路径\n",
    "        \"\"\"\n",
    "        # 加载嵌入\n",
    "        try:\n",
    "            self.text_embeddings = torch.load(text_embeddings_path)\n",
    "            self.image_embeddings = torch.load(image_embeddings_path)\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading embeddings: {e}\")\n",
    "            raise\n",
    "\n",
    "        # 确认嵌入数量匹配\n",
    "        assert len(self.text_embeddings) == len(self.image_embeddings), \"文本和图像嵌入的数量不匹配\"\n",
    "\n",
    "        # 确保嵌入的 dtype 为 float32\n",
    "        if self.text_embeddings.dtype != torch.float32:\n",
    "            self.text_embeddings = self.text_embeddings.float()\n",
    "            print(\"已将文本嵌入转换为 float32\")\n",
    "        if self.image_embeddings.dtype != torch.float32:\n",
    "            self.image_embeddings = self.image_embeddings.float()\n",
    "            print(\"已将图像嵌入转换为 float32\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.text_embeddings)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.text_embeddings[idx], self.image_embeddings[idx]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. 定义映射网络模型\n",
    "定义一个简单的全连接神经网络，将文本嵌入映射到图像嵌入空间。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义映射网络模型\n",
    "class TextToImageEmbedder(nn.Module):\n",
    "    def __init__(self, clip_dim=512, embed_dim=512):\n",
    "        super(TextToImageEmbedder, self).__init__()\n",
    "        self.mapping = nn.Sequential(\n",
    "            nn.Linear(clip_dim, 1024),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(1024, embed_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(embed_dim, embed_dim)\n",
    "        )\n",
    "    \n",
    "    def forward(self, text_embeddings):\n",
    "        image_embeddings = self.mapping(text_embeddings)\n",
    "        return image_embeddings\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. 创建数据加载器\n",
    "使用自定义的 Dataset 类创建 DataLoader，以便在训练过程中批量加载数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_6208/2941295110.py:11: 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",
      "  self.text_embeddings = torch.load(text_embeddings_path)\n",
      "/tmp/ipykernel_6208/2941295110.py:12: 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",
      "  self.image_embeddings = torch.load(image_embeddings_path)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已将文本嵌入转换为 float32\n",
      "已将图像嵌入转换为 float32\n",
      "数据集大小: 162770\n"
     ]
    }
   ],
   "source": [
    "# 创建映射网络的数据集和数据加载器（使用预计算嵌入）\n",
    "try:\n",
    "    mapping_dataset = CLIPEmbedMappingDatasetPrecomputed(\n",
    "        text_embeddings_path=TEXT_EMBEDDINGS_PATH,\n",
    "        image_embeddings_path=IMAGE_EMBEDDINGS_PATH\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(f\"Error initializing mapping dataset: {e}\")\n",
    "    raise\n",
    "\n",
    "# 确认数据集大小\n",
    "print(f\"数据集大小: {len(mapping_dataset)}\")\n",
    "\n",
    "# 创建 DataLoader\n",
    "mapping_loader = DataLoader(\n",
    "    mapping_dataset,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    num_workers=0,  # 根据需要调整\n",
    "    pin_memory=True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. 初始化模型、损失函数和优化器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建映射网络模型\n",
    "embedder = TextToImageEmbedder(clip_dim=512, embed_dim=512).to(device)\n",
    "\n",
    "# 定义损失函数和优化器\n",
    "criterion_mapping = nn.MSELoss()\n",
    "optimizer_mapping = optim.Adam(embedder.parameters(), lr=LEARNING_RATE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. 训练映射网络\n",
    "定义训练函数，并开始训练过程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 1/20: 100%|██████████| 1272/1272 [00:03<00:00, 413.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 1, 平均损失: 0.000563\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 2/20: 100%|██████████| 1272/1272 [00:02<00:00, 463.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 2, 平均损失: 0.000514\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 3/20: 100%|██████████| 1272/1272 [00:02<00:00, 471.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 3, 平均损失: 0.000507\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 4/20: 100%|██████████| 1272/1272 [00:02<00:00, 468.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 4, 平均损失: 0.000502\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 5/20: 100%|██████████| 1272/1272 [00:02<00:00, 490.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 5, 平均损失: 0.000500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 6/20: 100%|██████████| 1272/1272 [00:02<00:00, 438.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 6, 平均损失: 0.000498\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 7/20: 100%|██████████| 1272/1272 [00:02<00:00, 538.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 7, 平均损失: 0.000496\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 8/20: 100%|██████████| 1272/1272 [00:02<00:00, 491.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 8, 平均损失: 0.000495\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 9/20: 100%|██████████| 1272/1272 [00:02<00:00, 534.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 9, 平均损失: 0.000494\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 10/20: 100%|██████████| 1272/1272 [00:02<00:00, 543.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 10, 平均损失: 0.000493\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 11/20: 100%|██████████| 1272/1272 [00:02<00:00, 538.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 11, 平均损失: 0.000492\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 12/20: 100%|██████████| 1272/1272 [00:02<00:00, 473.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 12, 平均损失: 0.000491\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 13/20: 100%|██████████| 1272/1272 [00:02<00:00, 472.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 13, 平均损失: 0.000491\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 14/20: 100%|██████████| 1272/1272 [00:02<00:00, 481.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 14, 平均损失: 0.000490\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 15/20: 100%|██████████| 1272/1272 [00:02<00:00, 540.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 15, 平均损失: 0.000489\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 16/20: 100%|██████████| 1272/1272 [00:02<00:00, 506.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 16, 平均损失: 0.000489\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 17/20: 100%|██████████| 1272/1272 [00:02<00:00, 522.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 17, 平均损失: 0.000488\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 18/20: 100%|██████████| 1272/1272 [00:02<00:00, 556.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 18, 平均损失: 0.000488\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 19/20: 100%|██████████| 1272/1272 [00:02<00:00, 561.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 19, 平均损失: 0.000487\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 20/20: 100%|██████████| 1272/1272 [00:02<00:00, 499.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 20, 平均损失: 0.000487\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# 定义训练函数\n",
    "def train_mapping_network(model, loader, optimizer, criterion, num_epochs=20):\n",
    "    model.train()\n",
    "    for epoch in range(1, num_epochs + 1):\n",
    "        epoch_loss = 0\n",
    "        for text_emb, img_emb in tqdm(loader, desc=f\"Mapping Epoch {epoch}/{num_epochs}\"):\n",
    "            # 将嵌入移动到GPU\n",
    "            text_emb = text_emb.to(device)\n",
    "            img_emb = img_emb.to(device)\n",
    "\n",
    "            # 前向传播\n",
    "            pred_img_emb = model(text_emb)\n",
    "            loss = criterion(pred_img_emb, img_emb)\n",
    "\n",
    "            # 反向传播和优化\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "        avg_loss = epoch_loss / len(loader.dataset)\n",
    "        print(f\"Mapping Epoch {epoch}, 平均损失: {avg_loss:.6f}\")\n",
    "    return model\n",
    "\n",
    "# 开始训练\n",
    "trained_embedder = train_mapping_network(embedder, mapping_loader, optimizer_mapping, criterion_mapping, num_epochs=20)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "映射网络已保存为: text_to_image_embedder.pth\n"
     ]
    }
   ],
   "source": [
    "# 保存映射网络\n",
    "torch.save(trained_embedder.state_dict(), MODEL_SAVE_PATH)\n",
    "print(f\"映射网络已保存为: {MODEL_SAVE_PATH}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 利用后来生成的嵌入训练了一个新的网络效果比较好"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cuda\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_14605/393585030.py:40: 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",
      "  self.text_embeddings = torch.load(text_embeddings_path)\n",
      "/tmp/ipykernel_14605/393585030.py:41: 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",
      "  self.image_embeddings = torch.load(image_embeddings_path)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已将文本嵌入转换为 float32\n",
      "已将图像嵌入转换为 float32\n",
      "数据集大小: 162770\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 1/20: 100%|██████████| 1272/1272 [00:04<00:00, 301.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 1, 平均损失: 0.001742\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 2/20: 100%|██████████| 1272/1272 [00:04<00:00, 283.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 2, 平均损失: 0.001429\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 3/20: 100%|██████████| 1272/1272 [00:04<00:00, 288.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 3, 平均损失: 0.001401\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 4/20: 100%|██████████| 1272/1272 [00:04<00:00, 304.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 4, 平均损失: 0.001386\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 5/20: 100%|██████████| 1272/1272 [00:04<00:00, 298.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 5, 平均损失: 0.001377\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 6/20: 100%|██████████| 1272/1272 [00:04<00:00, 306.41it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 6, 平均损失: 0.001370\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 7/20: 100%|██████████| 1272/1272 [00:04<00:00, 295.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 7, 平均损失: 0.001365\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 8/20: 100%|██████████| 1272/1272 [00:04<00:00, 308.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 8, 平均损失: 0.001361\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 9/20: 100%|██████████| 1272/1272 [00:03<00:00, 355.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 9, 平均损失: 0.001357\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 10/20: 100%|██████████| 1272/1272 [00:03<00:00, 368.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 10, 平均损失: 0.001354\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 11/20: 100%|██████████| 1272/1272 [00:03<00:00, 321.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 11, 平均损失: 0.001351\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 12/20: 100%|██████████| 1272/1272 [00:04<00:00, 313.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 12, 平均损失: 0.001349\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 13/20: 100%|██████████| 1272/1272 [00:03<00:00, 322.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 13, 平均损失: 0.001347\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 14/20: 100%|██████████| 1272/1272 [00:03<00:00, 324.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 14, 平均损失: 0.001345\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 15/20: 100%|██████████| 1272/1272 [00:03<00:00, 348.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 15, 平均损失: 0.001342\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 16/20: 100%|██████████| 1272/1272 [00:03<00:00, 328.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 16, 平均损失: 0.001341\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 17/20: 100%|██████████| 1272/1272 [00:03<00:00, 329.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 17, 平均损失: 0.001339\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 18/20: 100%|██████████| 1272/1272 [00:04<00:00, 305.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 18, 平均损失: 0.001338\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 19/20: 100%|██████████| 1272/1272 [00:03<00:00, 323.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 19, 平均损失: 0.001336\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 20/20: 100%|██████████| 1272/1272 [00:03<00:00, 322.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 20, 平均损失: 0.001335\n",
      "映射网络已保存为: text_to_image_embedder.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# 导入必要的库\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "# 设置随机种子以确保可重复性\n",
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "random.seed(42)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(42)\n",
    "\n",
    "# 配置参数\n",
    "TEXT_EMBEDDINGS_PATH = '/root/autodl-tmp/clip_embeddings/text_embeddings_partition_0.pt'   # 替换为实际路径\n",
    "IMAGE_EMBEDDINGS_PATH = '/root/autodl-tmp/clip_embeddings/image_embeddings_partition_0.pt' # 替换为实际路径\n",
    "MODEL_SAVE_PATH = 'text_to_image_embedder.pth'\n",
    "BATCH_SIZE = 128\n",
    "NUM_EPOCHS = 20\n",
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "# 设备设置\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "# 定义自定义 Dataset 类\n",
    "class CLIPEmbedMappingDatasetPrecomputed(Dataset):\n",
    "    def __init__(self, text_embeddings_path, image_embeddings_path):\n",
    "        \"\"\"\n",
    "        初始化数据集，加载预计算的嵌入\n",
    "        :param text_embeddings_path: 文本嵌入文件路径\n",
    "        :param image_embeddings_path: 图像嵌入文件路径\n",
    "        \"\"\"\n",
    "        # 加载嵌入\n",
    "        try:\n",
    "            self.text_embeddings = torch.load(text_embeddings_path)\n",
    "            self.image_embeddings = torch.load(image_embeddings_path)\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading embeddings: {e}\")\n",
    "            raise\n",
    "\n",
    "        # 确认嵌入数量匹配\n",
    "        assert len(self.text_embeddings) == len(self.image_embeddings), \"文本和图像嵌入的数量不匹配\"\n",
    "\n",
    "        # 确保嵌入的 dtype 为 float32\n",
    "        if self.text_embeddings.dtype != torch.float32:\n",
    "            self.text_embeddings = self.text_embeddings.float()\n",
    "            print(\"已将文本嵌入转换为 float32\")\n",
    "        if self.image_embeddings.dtype != torch.float32:\n",
    "            self.image_embeddings = self.image_embeddings.float()\n",
    "            print(\"已将图像嵌入转换为 float32\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.text_embeddings)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.text_embeddings[idx], self.image_embeddings[idx]\n",
    "\n",
    "# 定义优化后的映射网络模型\n",
    "class TextToImageEmbedder(nn.Module):\n",
    "    def __init__(self, clip_dim=512, embed_dim=512):\n",
    "        super(TextToImageEmbedder, self).__init__()\n",
    "        self.mapping = nn.Sequential(\n",
    "            nn.Linear(clip_dim, 2048),\n",
    "            nn.BatchNorm1d(2048),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(2048, 1024),\n",
    "            nn.BatchNorm1d(1024),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(1024, embed_dim),\n",
    "            nn.BatchNorm1d(embed_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(embed_dim, embed_dim)\n",
    "        )\n",
    "    \n",
    "    def forward(self, text_embeddings):\n",
    "        image_embeddings = self.mapping(text_embeddings)\n",
    "        return image_embeddings\n",
    "\n",
    "# 定义新的余弦相似度损失函数\n",
    "class CosineSimilarityLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CosineSimilarityLoss, self).__init__()\n",
    "        self.cos = nn.CosineSimilarity(dim=1)\n",
    "    \n",
    "    def forward(self, pred, target):\n",
    "        # 余弦相似度范围在 [-1, 1]，我们希望最大化相似度，因此损失为 1 - cos_sim\n",
    "        return 1 - self.cos(pred, target).mean()\n",
    "\n",
    "# 创建映射网络的数据集和数据加载器（使用预计算嵌入）\n",
    "try:\n",
    "    mapping_dataset = CLIPEmbedMappingDatasetPrecomputed(\n",
    "        text_embeddings_path=TEXT_EMBEDDINGS_PATH,\n",
    "        image_embeddings_path=IMAGE_EMBEDDINGS_PATH\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(f\"Error initializing mapping dataset: {e}\")\n",
    "    raise\n",
    "\n",
    "# 确认数据集大小\n",
    "print(f\"数据集大小: {len(mapping_dataset)}\")\n",
    "\n",
    "# 创建 DataLoader\n",
    "mapping_loader = DataLoader(\n",
    "    mapping_dataset,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    num_workers=4,  # 根据需要调整\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "# 创建映射网络模型\n",
    "embedder = TextToImageEmbedder(clip_dim=512, embed_dim=512).to(device)\n",
    "\n",
    "# 定义损失函数和优化器\n",
    "criterion_mapping = CosineSimilarityLoss()\n",
    "optimizer_mapping = optim.Adam(embedder.parameters(), lr=LEARNING_RATE)\n",
    "\n",
    "# 定义训练函数\n",
    "def train_mapping_network(model, loader, optimizer, criterion, num_epochs=20):\n",
    "    model.train()\n",
    "    for epoch in range(1, num_epochs + 1):\n",
    "        epoch_loss = 0\n",
    "        for text_emb, img_emb in tqdm(loader, desc=f\"Mapping Epoch {epoch}/{num_epochs}\"):\n",
    "            # 将嵌入移动到GPU\n",
    "            text_emb = text_emb.to(device)\n",
    "            img_emb = img_emb.to(device)\n",
    "\n",
    "            # 前向传播\n",
    "            pred_img_emb = model(text_emb)\n",
    "            loss = criterion(pred_img_emb, img_emb)\n",
    "\n",
    "            # 反向传播和优化\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "        avg_loss = epoch_loss / len(loader.dataset)\n",
    "        print(f\"Mapping Epoch {epoch}, 平均损失: {avg_loss:.6f}\")\n",
    "    return model\n",
    "\n",
    "# 开始训练\n",
    "trained_embedder = train_mapping_network(embedder, mapping_loader, optimizer_mapping, criterion_mapping, num_epochs=NUM_EPOCHS)\n",
    "\n",
    "# 保存映射网络\n",
    "torch.save(trained_embedder.state_dict(), MODEL_SAVE_PATH)\n",
    "print(f\"映射网络已保存为: {MODEL_SAVE_PATH}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 下面这个映射函数太烂了 训练出来的是个屎"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cuda\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_14605/3116050070.py:40: 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",
      "  self.text_embeddings = torch.load(text_embeddings_path)\n",
      "/tmp/ipykernel_14605/3116050070.py:41: 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",
      "  self.image_embeddings = torch.load(image_embeddings_path)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已将文本嵌入转换为 float32\n",
      "已将图像嵌入转换为 float32\n",
      "数据集大小: 162770\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 1/20: 100%|██████████| 1272/1272 [00:09<00:00, 140.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 1, 平均训练损失: 0.002054\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 2/20: 100%|██████████| 1272/1272 [00:09<00:00, 132.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 2, 平均训练损失: 0.001955\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 3/20: 100%|██████████| 1272/1272 [00:10<00:00, 125.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 3, 平均训练损失: 0.001939\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 4/20: 100%|██████████| 1272/1272 [00:09<00:00, 131.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 4, 平均训练损失: 0.001937\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 5/20: 100%|██████████| 1272/1272 [00:09<00:00, 128.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 5, 平均训练损失: 0.001928\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 6/20: 100%|██████████| 1272/1272 [00:08<00:00, 141.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 6, 平均训练损失: 0.001924\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 7/20: 100%|██████████| 1272/1272 [00:08<00:00, 144.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 7, 平均训练损失: 0.001915\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 8/20: 100%|██████████| 1272/1272 [00:08<00:00, 141.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 8, 平均训练损失: 0.001917\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 9/20: 100%|██████████| 1272/1272 [00:09<00:00, 136.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 9, 平均训练损失: 0.001910\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 10/20: 100%|██████████| 1272/1272 [00:09<00:00, 139.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 10, 平均训练损失: 0.001909\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 11/20: 100%|██████████| 1272/1272 [00:08<00:00, 142.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 11, 平均训练损失: 0.001909\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 12/20: 100%|██████████| 1272/1272 [00:09<00:00, 140.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 12, 平均训练损失: 0.001902\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 13/20: 100%|██████████| 1272/1272 [00:09<00:00, 138.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 13, 平均训练损失: 0.001902\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 14/20: 100%|██████████| 1272/1272 [00:09<00:00, 136.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 14, 平均训练损失: 0.001904\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 15/20: 100%|██████████| 1272/1272 [00:09<00:00, 132.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 15, 平均训练损失: 0.001900\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 16/20: 100%|██████████| 1272/1272 [00:09<00:00, 140.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 16, 平均训练损失: 0.001895\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 17/20: 100%|██████████| 1272/1272 [00:09<00:00, 136.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 17, 平均训练损失: 0.001901\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 18/20: 100%|██████████| 1272/1272 [00:09<00:00, 133.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 18, 平均训练损失: 0.001895\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 19/20: 100%|██████████| 1272/1272 [00:09<00:00, 139.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 19, 平均训练损失: 0.001898\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 20/20: 100%|██████████| 1272/1272 [00:08<00:00, 148.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mapping Epoch 20, 平均训练损失: 0.001895\n",
      "映射网络已保存为: text_to_image_embedder.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# 导入必要的库\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "# 设置随机种子以确保可重复性\n",
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "random.seed(42)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(42)\n",
    "\n",
    "# 配置参数\n",
    "TEXT_EMBEDDINGS_PATH = '/root/autodl-tmp/clip_embeddings/text_embeddings_partition_0.pt'   # 替换为实际路径\n",
    "IMAGE_EMBEDDINGS_PATH = '/root/autodl-tmp/clip_embeddings/image_embeddings_partition_0.pt' # 替换为实际路径\n",
    "MODEL_SAVE_PATH = 'text_to_image_embedder.pth'\n",
    "BATCH_SIZE = 128\n",
    "NUM_EPOCHS = 20\n",
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "# 设备设置\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "# 定义自定义 Dataset 类\n",
    "class CLIPEmbedMappingDatasetPrecomputed(Dataset):\n",
    "    def __init__(self, text_embeddings_path, image_embeddings_path):\n",
    "        \"\"\"\n",
    "        初始化数据集，加载预计算的嵌入\n",
    "        :param text_embeddings_path: 文本嵌入文件路径\n",
    "        :param image_embeddings_path: 图像嵌入文件路径\n",
    "        \"\"\"\n",
    "        # 加载嵌入\n",
    "        try:\n",
    "            self.text_embeddings = torch.load(text_embeddings_path)\n",
    "            self.image_embeddings = torch.load(image_embeddings_path)\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading embeddings: {e}\")\n",
    "            raise\n",
    "\n",
    "        # 确认嵌入数量匹配\n",
    "        assert len(self.text_embeddings) == len(self.image_embeddings), \"文本和图像嵌入的数量不匹配\"\n",
    "\n",
    "        # 确保嵌入的 dtype 为 float32\n",
    "        if self.text_embeddings.dtype != torch.float32:\n",
    "            self.text_embeddings = self.text_embeddings.float()\n",
    "            print(\"已将文本嵌入转换为 float32\")\n",
    "        if self.image_embeddings.dtype != torch.float32:\n",
    "            self.image_embeddings = self.image_embeddings.float()\n",
    "            print(\"已将图像嵌入转换为 float32\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.text_embeddings)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.text_embeddings[idx], self.image_embeddings[idx]\n",
    "\n",
    "# 使用GELU或SiLU(Swish)激活函数\n",
    "# PyTorch中SiLU基本等同于Swish: y = x * sigmoid(x)\n",
    "# 我们这里使用nn.SiLU()来代替GELU，如果需要GELU可以将nn.SiLU()替换为nn.GELU()\n",
    "activation_fn = nn.SiLU()\n",
    "\n",
    "# 定义改进的映射网络模型\n",
    "class TextToImageEmbedder(nn.Module):\n",
    "    def __init__(self, clip_dim=512, embed_dim=512):\n",
    "        super(TextToImageEmbedder, self).__init__()\n",
    "        # 增加层数和宽度，并使用SiLU激活函数来替换ReLU\n",
    "        self.mapping = nn.Sequential(\n",
    "            nn.Linear(clip_dim, 4096),\n",
    "            nn.BatchNorm1d(4096),\n",
    "            activation_fn,\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(4096, 2048),\n",
    "            nn.BatchNorm1d(2048),\n",
    "            activation_fn,\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(2048, 1024),\n",
    "            nn.BatchNorm1d(1024),\n",
    "            activation_fn,\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(1024, embed_dim),\n",
    "            nn.BatchNorm1d(embed_dim),\n",
    "            activation_fn,\n",
    "            nn.Linear(embed_dim, embed_dim) # 输出层\n",
    "        )\n",
    "\n",
    "    def forward(self, text_embeddings):\n",
    "        image_embeddings = self.mapping(text_embeddings)\n",
    "        # 可选：对映射后的图像嵌入进行归一化，使其分布更稳定\n",
    "        image_embeddings = image_embeddings / (image_embeddings.norm(dim=1, keepdim=True) + 1e-8)\n",
    "        return image_embeddings\n",
    "\n",
    "# 定义对比损失 - 使用三元组损失作为简单的对比学习示例\n",
    "# anchor: 映射网络预测的图像嵌入 pred_img_emb\n",
    "# positive: 对应的真实图像嵌入 img_emb[i]\n",
    "# negative: 来自同一batch中其他样本的img_emb\n",
    "triplet_loss_fn = nn.TripletMarginLoss(margin=0.2)\n",
    "\n",
    "def contrastive_loss(pred_img_emb, img_emb):\n",
    "    \"\"\"\n",
    "    对比损失函数(简易版)：使用三元组损失进行对比学习\n",
    "    对于每个样本(i):\n",
    "      anchor = pred_img_emb[i]\n",
    "      positive = img_emb[i]\n",
    "      negative = img_emb[j], j != i\n",
    "\n",
    "    需要保证batch_size > 1\n",
    "    \"\"\"\n",
    "    batch_size = pred_img_emb.size(0)\n",
    "    # 随机选择negative样本\n",
    "    # 为了简单起见，从同一batch中随机选取与i不同的样本作为negative\n",
    "    indices = torch.arange(batch_size, device=pred_img_emb.device)\n",
    "    # 打乱indices\n",
    "    shuffled = indices[torch.randperm(batch_size)]\n",
    "    # 确保没有出现与i相同的索引\n",
    "    # 如果出现相同索引则再次随机\n",
    "    # 简化处理：若有相同索引则重排一次\n",
    "    for i in range(batch_size):\n",
    "        if shuffled[i] == i:\n",
    "            # 简单处理：将该位置与下一个位置交换\n",
    "            if i < batch_size - 1:\n",
    "                shuffled[i], shuffled[i+1] = shuffled[i+1], shuffled[i]\n",
    "            else:\n",
    "                shuffled[i], shuffled[i-1] = shuffled[i-1], shuffled[i]\n",
    "\n",
    "    negative_img_emb = img_emb[shuffled]\n",
    "    loss = triplet_loss_fn(pred_img_emb, img_emb, negative_img_emb)\n",
    "    return loss\n",
    "\n",
    "# 创建映射网络的数据集和数据加载器（使用预计算嵌入）\n",
    "try:\n",
    "    mapping_dataset = CLIPEmbedMappingDatasetPrecomputed(\n",
    "        text_embeddings_path=TEXT_EMBEDDINGS_PATH,\n",
    "        image_embeddings_path=IMAGE_EMBEDDINGS_PATH\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(f\"Error initializing mapping dataset: {e}\")\n",
    "    raise\n",
    "\n",
    "# 确认数据集大小\n",
    "print(f\"数据集大小: {len(mapping_dataset)}\")\n",
    "\n",
    "# 创建 DataLoader\n",
    "mapping_loader = DataLoader(\n",
    "    mapping_dataset,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    num_workers=4,  # 根据需要调整\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "# 创建映射网络模型\n",
    "embedder = TextToImageEmbedder(clip_dim=512, embed_dim=512).to(device)\n",
    "\n",
    "# 定义优化器\n",
    "optimizer_mapping = optim.Adam(embedder.parameters(), lr=LEARNING_RATE)\n",
    "\n",
    "# 定义训练函数\n",
    "def train_mapping_network(model, loader, optimizer, num_epochs=NUM_EPOCHS):\n",
    "    model.train()\n",
    "    for epoch in range(1, num_epochs + 1):\n",
    "        epoch_loss = 0\n",
    "        for text_emb, img_emb in tqdm(loader, desc=f\"Mapping Epoch {epoch}/{num_epochs}\"):\n",
    "            # 将嵌入移动到GPU\n",
    "            text_emb = text_emb.to(device)\n",
    "            img_emb = img_emb.to(device)\n",
    "\n",
    "            # 前向传播\n",
    "            pred_img_emb = model(text_emb)\n",
    "\n",
    "            # 使用对比损失函数(三元组损失)\n",
    "            loss = contrastive_loss(pred_img_emb, img_emb)\n",
    "\n",
    "            # 反向传播和优化\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "\n",
    "        avg_loss = epoch_loss / len(loader.dataset)\n",
    "        print(f\"Mapping Epoch {epoch}, 平均训练损失: {avg_loss:.6f}\")\n",
    "    return model\n",
    "\n",
    "# 开始训练\n",
    "trained_embedder = train_mapping_network(embedder, mapping_loader, optimizer_mapping, num_epochs=NUM_EPOCHS)\n",
    "\n",
    "# 保存映射网络\n",
    "torch.save(trained_embedder.state_dict(), MODEL_SAVE_PATH)\n",
    "print(f\"映射网络已保存为: {MODEL_SAVE_PATH}\")\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": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}