make code more general for larger models

This commit is contained in:
rasbt 2024-05-05 10:18:40 -05:00
parent 3328b29521
commit a63b0f626c

View File

@ -1316,7 +1316,8 @@
"metadata": {}, "metadata": {},
"source": [ "source": [
"- Then, we replace the output layer (`model.out_head`), which originally maps the layer inputs to 50,257 dimensions (the size of the vocabulary)\n", "- Then, we replace the output layer (`model.out_head`), which originally maps the layer inputs to 50,257 dimensions (the size of the vocabulary)\n",
"- Since we finetune the model for binary classification (predicting 2 classes, \"spam\" and \"ham\"), we can replace the output layer as shown below, which will be trainable by default" "- Since we finetune the model for binary classification (predicting 2 classes, \"spam\" and \"ham\"), we can replace the output layer as shown below, which will be trainable by default\n",
"- Note that we use `BASE_CONFIG[\"emb_dim\"]` (which is equal to 768 in the `\"gpt2-small (124M)\"` model) to keep the code below more general"
] ]
}, },
{ {
@ -1329,7 +1330,7 @@
"torch.manual_seed(123)\n", "torch.manual_seed(123)\n",
"\n", "\n",
"num_classes = 2\n", "num_classes = 2\n",
"model.out_head = torch.nn.Linear(in_features=768, out_features=num_classes)" "model.out_head = torch.nn.Linear(in_features=BASE_CONFIG[\"emb_dim\"], out_features=num_classes)"
] ]
}, },
{ {