Router_Train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d656b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import torch\n",
    "import os\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import shortuuid\n",
    "import sys\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "from moellava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN\n",
    "from moellava.conversation import conv_templates, SeparatorStyle\n",
    "from moellava.model.builder import load_pretrained_model\n",
    "from moellava.utils import disable_torch_init\n",
    "from moellava.mm_utils import tokenizer_image_token, process_images, get_model_name_from_path\n",
    "from moellava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria\n",
    "\n",
    "from PIL import Image\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34e55227",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = (\"./checkpoints-clip336/llavastablelm-1.6b-finetune\")\n",
    "model_name = get_model_name_from_path(model_path)\n",
    "tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name,device_map='auto')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41753e9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import torch\n",
    "from PIL import Image\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Load your json with modality\n",
    "\n",
    "with open('Your json files with modality for training.', 'r', encoding='utf-8') as file:\n",
    "    data = json.load(file)\n",
    "\n",
    "\n",
    "unique_modalities = set(item['modality'] for item in data)#make sure your json file with modality\n",
    "label_map = {modality: idx for idx, modality in enumerate(unique_modalities)}\n",
    "#label_map = {'MRI': 0, 'CT': 1, 'pathology': 2, 'X-Ray': 3}\n",
    "\n",
    "class PVQADataset(Dataset):\n",
    "    def __init__(self, data, tokenizer, image_processor, max_length, label_map):\n",
    "        self.data = data\n",
    "        self.tokenizer = tokenizer\n",
    "        self.image_processor = image_processor\n",
    "        self.max_length = max_length\n",
    "        self.label_map = label_map\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        item = self.data[idx]\n",
    "        qs = item['conversations'][0]['value']  # question\n",
    "        qs = qs.replace('<image>\\n', '')\n",
    "        image_file = \"image_path\" + item['image']  # image_path\n",
    "        label = self.label_map[item['modality']]  \n",
    "        \n",
    "        image = Image.open(image_file).convert('RGB')\n",
    "        images = self.image_processor['image'].preprocess(image, return_tensors='pt')['pixel_values']\n",
    "        \n",
    "        if getattr(model.config, 'mm_use_im_start_end', False):\n",
    "            qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\\n' + qs\n",
    "        else:\n",
    "            qs = DEFAULT_IMAGE_TOKEN + '\\n' + qs\n",
    "\n",
    "        conv = conv_templates['stablelm'].copy()\n",
    "        conv.append_message(conv.roles[0], qs)\n",
    "        conv.append_message(conv.roles[1], None)\n",
    "        prompt = conv.get_prompt()\n",
    "        input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt')\n",
    "\n",
    "        return input_ids.squeeze(0), images.squeeze(0), label\n",
    "\n",
    "max_length = 2048\n",
    "# Construct Dataset和DataLoader\n",
    "dataset = PVQADataset(data, tokenizer, image_processor, max_length, label_map)\n",
    "dataloader = DataLoader(dataset, batch_size=1, shuffle=True)\n",
    "\n",
    "# MLP model\n",
    "class MLPClassifier(nn.Module):\n",
    "    def __init__(self, input_dim, num_classes):\n",
    "        super(MLPClassifier, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_dim, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "embedding_dim = 2048  # Stablelm embedding dim==2048,phi2 as 4096\n",
    "num_classes = len(label_map)\n",
    "\n",
    "mlp = MLPClassifier(embedding_dim, num_classes).to(\"cuda\")\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(mlp.parameters(), lr=0.001)\n",
    "\n",
    "accumulation_steps = 32\n",
    "accumulated_embeddings = []\n",
    "accumulated_labels = []\n",
    "\n",
    "for epoch in range(10):  # 10epoch\n",
    "    for i, (input_ids, images, label) in enumerate(dataloader):\n",
    "        input_ids, images, label = input_ids.cuda(), images.to(\"cuda\", dtype=torch.float16), label.cuda()\n",
    "        \n",
    "        with torch.inference_mode():\n",
    "            outputs = model(input_ids, images=images, output_hidden_states=True, return_dict=True)\n",
    "            hidden_states = outputs.hidden_states \n",
    "\n",
    "        all_layer_mean = torch.mean(torch.stack(hidden_states), dim=0)  \n",
    "        embedding = all_layer_mean.mean(dim=1)  # Avg\n",
    "        accumulated_embeddings.append(embedding.float())\n",
    "        accumulated_labels.append(label)\n",
    "        \n",
    "        if (i + 1) % accumulation_steps == 0:\n",
    "            accumulated_embeddings = torch.cat(accumulated_embeddings, dim=0)\n",
    "            accumulated_labels = torch.cat(accumulated_labels, dim=0)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            outputs = mlp(accumulated_embeddings)\n",
    "            loss = criterion(outputs, accumulated_labels)\n",
    "            #print(f\"Loss: {loss.item()}\")\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            accumulated_embeddings = []\n",
    "            accumulated_labels = []\n",
    "\n",
    "    print(f\"Epoch {epoch + 1}, Loss: {loss.item()}\")\n",
    "\n",
    "print(\"Done\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff985250",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_save_path = 'save path'\n",
    "torch.save(mlp.state_dict(), model_save_path)\n",
    "print(f\"Model has been saved to {model_save_path}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}