BERT_Domain_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Znxt4pnXSL8E"
   },
   "source": [
    "**Domain Data preparation** \\\\\n",
    "The first step is to download the Amazon Product Review dataset from the UCSD repository and preprocess the data. \\\\\n",
    "\n",
    "We are using metadata and not the reviews for training, as we want to model to get familiarized with the different types of products and their uses. \\\\\n",
    "\n",
    "We have uploaded the dataset to our Google drive (can also use wget function to download it but it takes some time). In preprocessing, we will drop the columns we do not need. We are not editing the text as it is domain training. We will do more detailed preprocessing such as removing of unwanted characters, removal of stopwords, etc. in the task specific training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "caWP1kbgRGUZ"
   },
   "outputs": [],
   "source": [
    "import gzip\n",
    "import json \n",
    "import pandas as pd \n",
    "\n",
    "#Unzipping the data\n",
    "data = []\n",
    "count=0\n",
    "with gzip.open('meta_All_Beauty.json.gz') as f:\n",
    "  for l in f:\n",
    "    data.append(json.loads(l.strip()))\n",
    "#     count+=1\n",
    "#     if count>300000:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4QhqDhKputEF",
    "outputId": "2448c28c-c8b3-4ee3-e4de-778adecb24b8"
   },
   "outputs": [],
   "source": [
    "print(data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "07d_oT8buwBM"
   },
   "outputs": [],
   "source": [
    "#Loading the dataset into a Dataframe\n",
    "df = pd.DataFrame.from_dict(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "aDS1YWCNu8gj"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32892"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iEavdY5au9_t"
   },
   "outputs": [],
   "source": [
    "print(df.description.iloc[0], df.title.iloc[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LaCkVN4IyEMf"
   },
   "outputs": [],
   "source": [
    "df.description.iloc[1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "323e7ONpzgk-"
   },
   "outputs": [],
   "source": [
    "#Function to edit the description column of the dataframe \n",
    "def edit_description(desc):\n",
    "  if desc==[]:\n",
    "    return ''\n",
    "  else:\n",
    "    return desc[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "R3Ly3pu90E9f"
   },
   "outputs": [],
   "source": [
    "df['description'] = df.apply(lambda row: edit_description(row['description']), axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "Y5uvJNncvjrY"
   },
   "outputs": [],
   "source": [
    "#Filtering the dataset \n",
    "# 1) Check which all rows contain NA entries for title and description\n",
    "# 2) Split the dataframe into 2 halves based on the above condition \n",
    "df = df.fillna('')\n",
    "filtered_df = df[((df.title !='') & (df.description!=''))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Miw38NaR0a6r"
   },
   "outputs": [],
   "source": [
    "len(filtered_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "Ha8FZobx0csY"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\anshj\\AppData\\Local\\Temp\\ipykernel_2456\\2981561231.py:2: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  filtered_df['custom_input'] = filtered_df['title'] + filtered_df['description']\n"
     ]
    }
   ],
   "source": [
    "#Combining the title and description into one column - 'custom_input'\n",
    "filtered_df['custom_input'] = filtered_df['title'] + filtered_df['description']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "r0CIcqOM0sM1"
   },
   "outputs": [],
   "source": [
    "train_df = filtered_df[['custom_input']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HNR4rvCR2Y4j"
   },
   "outputs": [],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B2S0dJJd57jW"
   },
   "source": [
    "**Data Visualization** \\\\\n",
    "Will be checking the wordcloud to visually see which word occured the most in the training data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 198
    },
    "id": "ezOoY84c2an0",
    "outputId": "b5b05092-d4a9-432b-b57c-ddb532f2d8d8"
   },
   "outputs": [],
   "source": [
    "import nltk\n",
    "# nltk.download('stopwords')\n",
    "from nltk.corpus import stopwords\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from wordcloud import WordCloud \n",
    "\n",
    "stopwords = stopwords.words('english')\n",
    "stopwords = stopwords + [\"br\",\"href\"]\n",
    "text = \" \".join(input for input in train_df.custom_input)\n",
    "\n",
    "wordcloud = WordCloud(stopwords=stopwords).generate(text)\n",
    "\n",
    "plt.imshow(wordcloud, interpolation='bilinear')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qyDq_JFbiWiL"
   },
   "source": [
    "**Data Splitting** \\\\\n",
    "Splitting the data into train and validation sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OJto3-EdgwD9"
   },
   "outputs": [],
   "source": [
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "DoT8Q7VJ-vcK"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\anshj\\anaconda3\\envs\\torch-gpu\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch \n",
    "from transformers import BertTokenizer, BertForMaskedLM\n",
    "from torch.utils.data import DataLoader, Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "gdgb9BkmhXBs"
   },
   "outputs": [],
   "source": [
    "#Data split - 80% training and 20% validaiton \n",
    "train_data = train_df.sample(frac=1, random_state=42)\n",
    "train_data = train_data.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "rctkf0IgfLDU"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "#Initializing the tokenizer and model\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "model = BertForMaskedLM.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gln0o99cicWu"
   },
   "source": [
    "**Tokenization** \\\\\n",
    "Tokenize the text data using the BERT tokenizer. This involves converting the text into a sequence of tokens that can be fed into the BERT model.\n",
    "\n",
    "Due to limited size of the dataset we are using the base model with 110M parameters and not the BERT large model which has 330M parameters. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ROQfyPEDrPqw"
   },
   "source": [
    "There are several other options of pretrained tokenizers, but we are using BERT Tokenizer. A tokenizer splits text into tokens according to a set of rules. The tokens are converted into numbers and then tensors, which become the model inputs. Finally, you want the tokenizer to return the actual tensors that get fed to the model. Set the return_tensors parameter to either pt for PyTorch, or tf for TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "Vc-QnDUzfW2E"
   },
   "outputs": [],
   "source": [
    "#DataLoader expects 2 additional fucntions apart from the initialized Dataset object - getitem and len methods\n",
    "class AmazonReviewsDataset(Dataset):\n",
    "  #This class handles the conversion of data into a Dataset object\n",
    "    def __init__(self, data, tokenizer):\n",
    "        self.data = data\n",
    "        self.tokenizer = tokenizer\n",
    "    def __getitem__(self, index): \n",
    "        title_desc = self.data.loc[index, 'custom_input']\n",
    "        input = self.tokenizer(title_desc, return_tensors='pt', add_special_tokens=True, max_length =512, truncation=True, padding='max_length')\n",
    "        input_ids = input.input_ids.squeeze()\n",
    "        attention_mask = input.attention_mask.squeeze()\n",
    "        labels = input_ids.clone()\n",
    "        #creates a copy of input ids as the label \n",
    "        \n",
    "        #Now we are creating the mask - each non-special token has a 15% chance of getting masked \n",
    "        rand = torch.rand(input_ids.shape)\n",
    "        #create a tensor of float values (b/w 0-1) that have equal dimensions as our input\n",
    "        mask_arr = (rand<0.15) * (input.input_ids[0]!=101) * (input.input_ids[0]!=102) * (input.input_ids[0]!=0)\n",
    "        selection = torch.flatten(mask_arr.nonzero()).tolist()\n",
    "        input_ids[selection] = 103\n",
    "        return {'input_ids':input_ids, 'attention_mask': attention_mask, 'labels': labels}\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "train_dataset = AmazonReviewsDataset(train_data, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4P4ADAhuhvDo"
   },
   "source": [
    "If we directly call the tokenizer instead of using the encode function, it will return the token_ids and attention mask as well. And therefore we do not have to create input_mask on our own. Moreover, as we have specified that we want the returned output to be a pytorch tensor, therefore no need to use torch.tensor() function while returning the tokenized input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Uj9BoFCvle2d"
   },
   "outputs": [],
   "source": [
    "train_dataset.__getitem__(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7HKZDTQWuYTP"
   },
   "source": [
    "**Model Domain Training** \\\\\n",
    "We are training model from the scratch for 20 epochs on the dataset in a  Masked Language Modeling fashion. BERT model can be trained in 2 possible ways - 1) MLM 2) NSP (Next Sequence Prediction). As we are using MLM only, therefore we are making use of BertMaskedLM model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "71vLqg5Hr0hC"
   },
   "source": [
    "[SEP] - 102 \\\\\n",
    "[MASK] - 103 \\\\\n",
    "[CLS] - 101 \\\\\n",
    "We are only conserned about the input_ids and won't give much attention to token_type_ids (useful only when we have more than one sequences) and attention_mask (tells whether we have to attend a token or not). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F21PSgifztwf",
    "outputId": "e5e38c84-e210-46fb-e64e-a89c97f3f191"
   },
   "outputs": [],
   "source": [
    "len(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "u-erwuBAik3g"
   },
   "outputs": [],
   "source": [
    "#We won't need to focus on the token_ids, because we are doing MLM\n",
    "#For training we need input_ids with the mask tokens and output labels (without the mask tokens)\n",
    "\n",
    "#Now we have our tokens in correct format and dimensioanlity but we need to process them through DataLoader object\n",
    "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "yVr_I6Ybnh-_"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForMaskedLM(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (cls): BertOnlyMLMHead(\n",
       "    (predictions): BertLMPredictionHead(\n",
       "      (transform): BertPredictionHeadTransform(\n",
       "        (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (transform_act_fn): GELUActivation()\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      )\n",
       "      (decoder): Linear(in_features=768, out_features=30522, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Checking if we have a GPU and if we have move the model to it \n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "MhyjG92pn8ws"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForMaskedLM(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (cls): BertOnlyMLMHead(\n",
       "    (predictions): BertLMPredictionHead(\n",
       "      (transform): BertPredictionHeadTransform(\n",
       "        (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (transform_act_fn): GELUActivation()\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      )\n",
       "      (decoder): Linear(in_features=768, out_features=30522, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Enabling model's training mode\n",
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4pRbLJ-0oB3B",
    "outputId": "3554665d-5be4-40fd-a5fa-4d41d7011d72"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\anshj\\anaconda3\\envs\\torch-gpu\\lib\\site-packages\\transformers\\optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from transformers import AdamW\n",
    "#We are using Adam with Weight Decay as the optimizer \n",
    "optim = AdamW(model.parameters(), lr=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "qx_mkahtoXgq"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0: 100%|█████████████████████████████████████████████████████████| 1823/1823 [07:17<00:00,  4.16it/s, loss=0.114]\n",
      "Epoch 1: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:17<00:00,  4.17it/s, loss=0.0455]\n",
      "Epoch 2: 100%|█████████████████████████████████████████████████████████| 1823/1823 [07:17<00:00,  4.17it/s, loss=0.146]\n",
      "Epoch 3: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:16<00:00,  4.17it/s, loss=0.0709]\n",
      "Epoch 4: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:16<00:00,  4.17it/s, loss=0.0573]\n",
      "Epoch 5: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:16<00:00,  4.17it/s, loss=0.0651]\n",
      "Epoch 6: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:16<00:00,  4.17it/s, loss=0.0493]\n",
      "Epoch 7: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:17<00:00,  4.17it/s, loss=0.0577]\n",
      "Epoch 8: 100%|████████████████████████████████████████████████████████| 1823/1823 [07:17<00:00,  4.17it/s, loss=0.0814]\n",
      "Epoch 9: 100%|█████████████████████████████████████████████████████████| 1823/1823 [07:16<00:00,  4.17it/s, loss=0.155]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "#tqdm allows us to create a progress bar during training \n",
    "\n",
    "epochs = 10\n",
    "#Be careful with number of epochs when training transformer models - they tend to overfit very easily \n",
    "\n",
    "for epoch in range(epochs):\n",
    "  #Start training loop \n",
    "  loop = tqdm(train_loader, leave = True)\n",
    "  #leave=True, leaves a progress bar rather than replacing it with a new one after each epoch\n",
    "  \n",
    "  #iterating over batches in loop\n",
    "  for batch in loop:\n",
    "    optim.zero_grad()\n",
    "    #Rather than having any randomly initialized gradients at the start, we make them zero \n",
    "    \n",
    "    input_ids = batch['input_ids'].to(device) \n",
    "    attention_mask = batch['attention_mask'].to(device)\n",
    "    labels = batch['labels'].to(device)\n",
    "    #We want the data as well as the model to be present in the GPU \n",
    "\n",
    "    outputs = model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "    #Extracting loss from those outputs\n",
    "\n",
    "    loss = outputs.loss \n",
    "    loss.backward() #Calculates loss for every parameter in our model, now can do the gradient update using the optimizer\n",
    "    optim.step() #Takes a step to optimize every parameter in our model\n",
    "\n",
    "    loop.set_description(f'Epoch {epoch}') \n",
    "    loop.set_postfix(loss=loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "ht8dUS8LvhuZ"
   },
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'Model_Domain_weigths')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python (torch-gpu)",
   "language": "python",
   "name": "torch-gpu"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}