<h1>Fine Tuning Approach</h1>
<p>There are multiple approaches to fine-tune BERT for the target tasks.</p>
<ol>
<li>Further Pre-training the base BERT model</li>
<li>Custom classification layer(s) on top of the base BERT model  being trainable</li>
<li>Custom classification layer(s) on top of the base BERT model being non-trainable (frozen)</li>
</ol>
<p>Note that the BERT base model has been pre-trained only for two tasks as in the original paper.</p>
<ul>
<li><a href="https://arxiv.org/abs/1810.04805" rel="noreferrer">BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding</a></li>
</ul>
<blockquote>
<p>3.1 Pre-training BERT ...we pre-train BERT using two unsupervised tasks<br></p>
<ul>
<li>Task #1: Masked LM<br></li>
<li>Task #2: Next Sentence Prediction (NSP)<br></li>
</ul>
</blockquote>
<p>Hence, the base BERT model is like half-baked which can be fully baked for the target domain (1st way). We can use it as part of our custom model training with the base trainable (2nd) or not-trainable (3rd).</p>
<hr />
<h1>1st approach</h1>
<p><a href="https://arxiv.org/abs/1905.05583" rel="noreferrer">How to Fine-Tune BERT for Text Classification?</a> demonstrated the 1st approach of Further Pre-training, and pointed out the learning rate is the key to avoid <strong>Catastrophic Forgetting</strong> where the pre-trained knowledge is erased during learning of new knowledge.</p>
<blockquote>
<p>We find that a lower learning rate, such as 2e-5,
is necessary to make BERT overcome the catastrophic forgetting problem. With an aggressive learn rate of 4e-4, the training set fails to converge.<br>
<a href="https://i.stack.imgur.com/pm2EV.png" rel="noreferrer"><img src="https://i.stack.imgur.com/pm2EV.png" alt="enter image description here" /></a></p>
</blockquote>
<p>Probably this is the reason why the <a href="https://arxiv.org/abs/1810.04805" rel="noreferrer">BERT paper</a> used 5e-5, 4e-5, 3e-5, and 2e-5 for <strong>fine-tuning</strong>.</p>
<blockquote>
<p>We use a batch size of 32 and fine-tune for 3 epochs over the data for all GLUE tasks. For each task, we selected the best fine-tuning learning rate (among 5e-5, 4e-5, 3e-5, and 2e-5) on the Dev set</p>
</blockquote>
<p>Note that the base model pre-training itself used higher learning rate.</p>
<ul>
<li><a href="https://huggingface.co/bert-base-uncased#pretraining" rel="noreferrer">bert-base-uncased - pretraining</a></li>
</ul>
<blockquote>
<p>The model was trained on 4 cloud TPUs in Pod configuration (16 TPU chips total) for one million steps with a batch size of 256. The sequence length was limited to 128 tokens for 90% of the steps and 512 for the remaining 10%. The optimizer used is Adam with a learning rate of <code>1e-4</code>, β1=<code>0.9</code> and β2=<code>0.999</code>, a weight decay of <code>0.01</code>, learning rate warmup for 10,000 steps and linear decay of the learning rate after.</p>
</blockquote>
<p>Will describe the 1st way as part of the 3rd approach below.</p>
<p>FYI:
<a href="https://huggingface.co/transformers/model_doc/distilbert.html#tfdistilbertmodel" rel="noreferrer">TFDistilBertModel</a> is the bare base model with the name <code>distilbert</code>.</p>
<pre><code>Model: &quot;tf_distil_bert_model_1&quot;
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
distilbert (TFDistilBertMain multiple                  66362880  
=================================================================
Total params: 66,362,880
Trainable params: 66,362,880
Non-trainable params: 0
</code></pre>
<hr />
<h1>2nd approach</h1>
<p>Huggingface takes the 2nd approach as in <a href="https://huggingface.co/transformers/custom_datasets.html#fine-tuning-with-native-pytorch-tensorflow" rel="noreferrer">Fine-tuning with native PyTorch/TensorFlow</a> where <code>TFDistilBertForSequenceClassification</code> has added the custom classification layer <code>classifier</code> on top of the base <code>distilbert</code> model being trainable. The small learning rate requirement will apply as well to avoid the catastrophic forgetting.</p>
<pre><code>from transformers import TFDistilBertForSequenceClassification

model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
optimizer = tf.keras.optimizers.Adam(learning_rate=5e-5)
model.compile(optimizer=optimizer, loss=model.compute_loss) # can also use any keras loss fn
model.fit(train_dataset.shuffle(1000).batch(16), epochs=3, batch_size=16)
</code></pre>
<pre><code>Model: &quot;tf_distil_bert_for_sequence_classification_2&quot;
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
distilbert (TFDistilBertMain multiple                  66362880  
_________________________________________________________________
pre_classifier (Dense)       multiple                  590592    
_________________________________________________________________
classifier (Dense)           multiple                  1538      
_________________________________________________________________
dropout_59 (Dropout)         multiple                  0         
=================================================================
Total params: 66,955,010
Trainable params: 66,955,010  &lt;--- All parameters are trainable
Non-trainable params: 0
</code></pre>
<h3>Implementation of the 2nd approach</h3>
<pre><code>import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split
from transformers import (
    DistilBertTokenizerFast,
    TFDistilBertForSequenceClassification,
)


DATA_COLUMN = 'text'
LABEL_COLUMN = 'category_index'
MAX_SEQUENCE_LENGTH = 512
LEARNING_RATE = 5e-5
BATCH_SIZE = 16
NUM_EPOCHS = 3


# --------------------------------------------------------------------------------
# Tokenizer
# --------------------------------------------------------------------------------
tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
def tokenize(sentences, max_length=MAX_SEQUENCE_LENGTH, padding='max_length'):
    &quot;&quot;&quot;Tokenize using the Huggingface tokenizer
    Args:
        sentences: String or list of string to tokenize
        padding: Padding method ['do_not_pad'|'longest'|'max_length']
    &quot;&quot;&quot;
    return tokenizer(
        sentences,
        truncation=True,
        padding=padding,
        max_length=max_length,
        return_tensors=&quot;tf&quot;
    )

# --------------------------------------------------------------------------------
# Load data
# --------------------------------------------------------------------------------
raw_train = pd.read_csv(&quot;./train.csv&quot;)
train_data, validation_data, train_label, validation_label = train_test_split(
    raw_train[DATA_COLUMN].tolist(),
    raw_train[LABEL_COLUMN].tolist(),
    test_size=.2,
    shuffle=True
)

# --------------------------------------------------------------------------------
# Prepare TF dataset
# --------------------------------------------------------------------------------
train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(tokenize(train_data)),  # Convert BatchEncoding instance to dictionary
    train_label
)).shuffle(1000).batch(BATCH_SIZE).prefetch(1)
validation_dataset = tf.data.Dataset.from_tensor_slices((
    dict(tokenize(validation_data)),
    validation_label
)).batch(BATCH_SIZE).prefetch(1)

# --------------------------------------------------------------------------------
# training
# --------------------------------------------------------------------------------
model = TFDistilBertForSequenceClassification.from_pretrained(
    'distilbert-base-uncased',
    num_labels=NUM_LABELS
)
optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)
model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
model.fit(
    x=train_dataset,
    y=None,
    validation_data=validation_dataset,
    batch_size=BATCH_SIZE,
    epochs=NUM_EPOCHS,
)
</code></pre>
<hr />
<h1>3rd approach</h1>
<h2>Basics</h2>
<p>Please note that the images are taken from <a href="http://jalammar.github.io/a-visual-guide-to-using-bert-for-the-first-time/" rel="noreferrer">A Visual Guide to Using BERT for the First Time</a> and modified.</p>
<h3>Tokenizer</h3>
<p><a href="https://huggingface.co/transformers/main_classes/tokenizer.html" rel="noreferrer">Tokenizer</a> generates the instance of BatchEncoding which can be used like a Python dictionary and the input to the BERT model.</p>
<ul>
<li><a href="https://huggingface.co/transformers/main_classes/tokenizer.html#batchencoding" rel="noreferrer">BatchEncoding</a></li>
</ul>
<blockquote>
<p>Holds the output of the encode_plus() and batch_encode() methods (tokens, attention_masks, etc).
<br>
This class is derived from a python dictionary and <strong>can be used as a dictionary</strong>. In addition, this class exposes utility methods to map from word/character space to token space.<br><br>
Parameters<br></p>
<ul>
<li>data (dict) – Dictionary of lists/arrays/tensors returned by the encode/batch_encode methods (‘input_ids’, ‘attention_mask’, etc.).</li>
</ul>
</blockquote>
<p>The <code>data</code> attribute of the class is the tokens generated which has <code>input_ids</code> and <code>attention_mask</code> elements.</p>
<h3>input_ids</h3>
<ul>
<li><a href="https://huggingface.co/transformers/glossary.html#input-ids" rel="noreferrer">input_ids</a></li>
</ul>
<blockquote>
<p>The input ids are often the only required parameters to be passed to the model as input. They are <strong>token indices, numerical representations of tokens</strong> building the sequences that will be used as input by the model.</p>
</blockquote>
<h3>attention_mask</h3>
<ul>
<li><a href="https://huggingface.co/transformers/glossary.html#attention-mask" rel="noreferrer">Attention mask</a></li>
</ul>
<blockquote>
<p>This argument indicates to the model which tokens should be attended to, and which should not.</p>
</blockquote>
<p>If the attention_mask is <code>0</code>, the token id is ignored. For instance if a sequence is padded to adjust the sequence length, the padded words should be ignored hence their attention_mask are 0.</p>
<h3>Special Tokens</h3>
<p>BertTokenizer addes special tokens, enclosing a sequence with <code>[CLS]</code> and <code>[SEP]</code>. <code>[CLS]</code> represents <strong>Classification</strong> and <code>[SEP]</code> separates sequences. For Question Answer or Paraphrase tasks, <code>[SEP]</code> separates the two sentences to compare.</p>
<p><a href="https://huggingface.co/transformers/model_doc/bert.html#berttokenizer" rel="noreferrer">BertTokenizer</a></p>
<blockquote>
<ul>
<li>cls_token (str, optional, defaults to &quot;<strong>[CLS]</strong>&quot;)<BR>The <strong>Classifier Token which is used when doing sequence classification</strong> (classification of the whole sequence instead of per-token classification). It is the first token of the sequence when built with special tokens.</li>
<li>sep_token (str, optional, defaults to &quot;[SEP]&quot;)<BR>The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for sequence classification or for a text and a question for question answering. It is also used as the last token of a sequence built with special tokens.</li>
</ul>
</blockquote>
<p><a href="http://jalammar.github.io/a-visual-guide-to-using-bert-for-the-first-time/" rel="noreferrer">A Visual Guide to Using BERT for the First Time</a> show the tokenization.</p>
<p><a href="https://i.stack.imgur.com/zQtff.png" rel="noreferrer"><img src="https://i.stack.imgur.com/zQtff.png" alt="enter image description here" /></a></p>
<h3>[CLS]</h3>
<p>The embedding vector for <strong><code>[CLS]</code></strong> in the output from the base model final layer represents the classification that has been learned by the base model. Hence feed the embedding vector of  <strong><code>[CLS]</code></strong> token into the classification layer added on top of the base model.</p>
<ul>
<li><a href="https://arxiv.org/abs/1810.04805" rel="noreferrer">BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding</a></li>
</ul>
<blockquote>
<p>The first token of every sequence is always <code>a special classification token ([CLS])</code>. The final hidden state corresponding to this token is <strong>used as the aggregate sequence representation for classification tasks</strong>. Sentence pairs are packed together into a single sequence. We differentiate the sentences in two ways. First, we separate them with a special token ([SEP]). Second, we add a learned embedding to every token indicating whether it belongs to sentence A or sentence B.</p>
</blockquote>
<p>The model structure will be illustrated as below.</p>
<p><a href="https://i.stack.imgur.com/VAq7v.jpg" rel="noreferrer"><img src="https://i.stack.imgur.com/VAq7v.jpg" alt="enter image description here" /></a></p>
<p><a href="https://i.stack.imgur.com/tjpn4.jpg" rel="noreferrer"><img src="https://i.stack.imgur.com/tjpn4.jpg" alt="enter image description here" /></a></p>
<h3>Vector size</h3>
<p>In the model <code>distilbert-base-uncased</code>, each token is embedded into a vector of size <strong>768</strong>. The shape of the output from the base model is <code>(batch_size, max_sequence_length, embedding_vector_size=768)</code>. This accords with the BERT paper about the BERT/BASE model (as indicated in distilbert-<em><strong>base</strong></em>-uncased).</p>
<ul>
<li><a href="https://arxiv.org/abs/1810.04805" rel="noreferrer">BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding</a></li>
</ul>
<blockquote>
<p>BERT/BASE (L=12, H=<strong>768</strong>, A=12, Total Parameters=110M) and BERT/LARGE (L=24, H=1024, A=16, Total Parameters=340M).</p>
</blockquote>
<h3>Base Model - TFDistilBertModel</h3>
<ul>
<li><a href="https://towardsdatascience.com/hugging-face-transformers-fine-tuning-distilbert-for-binary-classification-tasks-490f1d192379" rel="noreferrer">Hugging Face Transformers: Fine-tuning DistilBERT for Binary Classification Tasks</a></li>
</ul>
<blockquote>
<p>TFDistilBertModel class to instantiate the base DistilBERT model <strong>without any specific head on top</strong> (as opposed to other classes such as TFDistilBertForSequenceClassification that do have an added classification head). <br><br>
We do not want any task-specific head attached because we simply want the pre-trained weights of the base model to provide a general understanding of the English language, and it will be our job to add our own classification head during the fine-tuning process in order to help the model distinguish between toxic comments.</p>
</blockquote>
<p><code>TFDistilBertModel</code> generates an instance of <code>TFBaseModelOutput</code> whose <code>last_hidden_state</code> parameter is the output from the model last layer.</p>
<pre><code>TFBaseModelOutput([(
    'last_hidden_state',
    &lt;tf.Tensor: shape=(batch_size, sequence_lendgth, 768), dtype=float32, numpy=array([[[...]]], dtype=float32)&gt;
)])
</code></pre>
<ul>
<li><a href="https://huggingface.co/transformers/main_classes/output.html#tfbasemodeloutput" rel="noreferrer">TFBaseModelOutput</a></li>
</ul>
<blockquote>
<p>Parameters<br></p>
<ul>
<li>last_hidden_state (tf.Tensor of shape (batch_size, sequence_length, hidden_size)) – Sequence of hidden-states at the output of the last layer of the model.</li>
</ul>
</blockquote>
<h2>Implementation</h2>
<h3>Python modules</h3>
<pre><code>import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split
from transformers import (
    DistilBertTokenizerFast,
    TFDistilBertModel,
)
</code></pre>
<h3>Configuration</h3>
<pre><code>TIMESTAMP = datetime.datetime.now().strftime(&quot;%Y%b%d%H%M&quot;).upper()

DATA_COLUMN = 'text'
LABEL_COLUMN = 'category_index'

MAX_SEQUENCE_LENGTH = 512   # Max length allowed for BERT is 512.
NUM_LABELS = len(raw_train[LABEL_COLUMN].unique())

MODEL_NAME = 'distilbert-base-uncased'
NUM_BASE_MODEL_OUTPUT = 768

# Flag to freeze base model
FREEZE_BASE = True

# Flag to add custom classification heads
USE_CUSTOM_HEAD = True
if USE_CUSTOM_HEAD == False:
    # Make the base trainable when no classification head exists.
    FREEZE_BASE = False


BATCH_SIZE = 16
LEARNING_RATE = 1e-2 if FREEZE_BASE else 5e-5
L2 = 0.01
</code></pre>
<h3>Tokenizer</h3>
<pre><code>tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_NAME)
def tokenize(sentences, max_length=MAX_SEQUENCE_LENGTH, padding='max_length'):
    &quot;&quot;&quot;Tokenize using the Huggingface tokenizer
    Args:
        sentences: String or list of string to tokenize
        padding: Padding method ['do_not_pad'|'longest'|'max_length']
    &quot;&quot;&quot;
    return tokenizer(
        sentences,
        truncation=True,
        padding=padding,
        max_length=max_length,
        return_tensors=&quot;tf&quot;
    )
</code></pre>
<h3>Input layer</h3>
<p>The base model expects <code>input_ids</code> and <code>attention_mask</code> whose shape is <code>(max_sequence_length,)</code>. Generate Keras Tensors for them with <code>Input</code> layer respectively.</p>
<pre><code># Inputs for token indices and attention masks
input_ids = tf.keras.layers.Input(shape=(MAX_SEQUENCE_LENGTH,), dtype=tf.int32, name='input_ids')
attention_mask = tf.keras.layers.Input((MAX_SEQUENCE_LENGTH,), dtype=tf.int32, name='attention_mask')
</code></pre>
<h3>Base model layer</h3>
<p>Generate the output from the base model. The base model generates <code>TFBaseModelOutput</code>. Feed the embedding of <strong><code>[CLS]</code></strong> to the next layer.</p>
<pre><code>base = TFDistilBertModel.from_pretrained(
    MODEL_NAME,
    num_labels=NUM_LABELS
)

# Freeze the base model weights.
if FREEZE_BASE:
    for layer in base.layers:
        layer.trainable = False
    base.summary()

# [CLS] embedding is last_hidden_state[:, 0, :]
output = base([input_ids, attention_mask]).last_hidden_state[:, 0, :]
</code></pre>
<h3>Classification layers</h3>
<pre><code>if USE_CUSTOM_HEAD:
    # -------------------------------------------------------------------------------
    # Classifiation leayer 01
    # --------------------------------------------------------------------------------
    output = tf.keras.layers.Dropout(
        rate=0.15,
        name=&quot;01_dropout&quot;,
    )(output)
    
    output = tf.keras.layers.Dense(
        units=NUM_BASE_MODEL_OUTPUT,
        kernel_initializer='glorot_uniform',
        activation=None,
        name=&quot;01_dense_relu_no_regularizer&quot;,
    )(output)
    output = tf.keras.layers.BatchNormalization(
        name=&quot;01_bn&quot;
    )(output)
    output = tf.keras.layers.Activation(
        &quot;relu&quot;,
        name=&quot;01_relu&quot;
    )(output)

    # --------------------------------------------------------------------------------
    # Classifiation leayer 02
    # --------------------------------------------------------------------------------
    output = tf.keras.layers.Dense(
        units=NUM_BASE_MODEL_OUTPUT,
        kernel_initializer='glorot_uniform',
        activation=None,
        name=&quot;02_dense_relu_no_regularizer&quot;,
    )(output)
    output = tf.keras.layers.BatchNormalization(
        name=&quot;02_bn&quot;
    )(output)
    output = tf.keras.layers.Activation(
        &quot;relu&quot;,
        name=&quot;02_relu&quot;
    )(output)
</code></pre>
<h3>Softmax Layer</h3>
<pre><code>output = tf.keras.layers.Dense(
    units=NUM_LABELS,
    kernel_initializer='glorot_uniform',
    kernel_regularizer=tf.keras.regularizers.l2(l2=L2),
    activation='softmax',
    name=&quot;softmax&quot;
)(output)
</code></pre>
<h3>Final Custom Model</h3>
<pre><code>name = f&quot;{TIMESTAMP}_{MODEL_NAME.upper()}&quot;
model = tf.keras.models.Model(inputs=[input_ids, attention_mask], outputs=output, name=name)
model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),
    metrics=['accuracy']
)
model.summary()
---
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_ids (InputLayer)          [(None, 256)]        0                                            
__________________________________________________________________________________________________
attention_mask (InputLayer)     [(None, 256)]        0                                            
__________________________________________________________________________________________________
tf_distil_bert_model (TFDistilB TFBaseModelOutput(la 66362880    input_ids[0][0]                  
                                                                 attention_mask[0][0]             
__________________________________________________________________________________________________
tf.__operators__.getitem_1 (Sli (None, 768)          0           tf_distil_bert_model[1][0]       
__________________________________________________________________________________________________
01_dropout (Dropout)            (None, 768)          0           tf.__operators__.getitem_1[0][0] 
__________________________________________________________________________________________________
01_dense_relu_no_regularizer (D (None, 768)          590592      01_dropout[0][0]                 
__________________________________________________________________________________________________
01_bn (BatchNormalization)      (None, 768)          3072        01_dense_relu_no_regularizer[0][0
__________________________________________________________________________________________________
01_relu (Activation)            (None, 768)          0           01_bn[0][0]                      
__________________________________________________________________________________________________
02_dense_relu_no_regularizer (D (None, 768)          590592      01_relu[0][0]                    
__________________________________________________________________________________________________
02_bn (BatchNormalization)      (None, 768)          3072        02_dense_relu_no_regularizer[0][0
__________________________________________________________________________________________________
02_relu (Activation)            (None, 768)          0           02_bn[0][0]                      
__________________________________________________________________________________________________
softmax (Dense)                 (None, 2)            1538        02_relu[0][0]                    
==================================================================================================
Total params: 67,551,746
Trainable params: 1,185,794
Non-trainable params: 66,365,952   &lt;--- Base BERT model is frozen
</code></pre>
<h3>Data allocation</h3>
<pre><code># --------------------------------------------------------------------------------
# Split data into training and validation
# --------------------------------------------------------------------------------
raw_train = pd.read_csv(&quot;./train.csv&quot;)
train_data, validation_data, train_label, validation_label = train_test_split(
    raw_train[DATA_COLUMN].tolist(),
    raw_train[LABEL_COLUMN].tolist(),
    test_size=.2,
    shuffle=True
)

# X = dict(tokenize(train_data))
# Y = tf.convert_to_tensor(train_label)
X = tf.data.Dataset.from_tensor_slices((
    dict(tokenize(train_data)),  # Convert BatchEncoding instance to dictionary
    train_label
)).batch(BATCH_SIZE).prefetch(1)

V = tf.data.Dataset.from_tensor_slices((
    dict(tokenize(validation_data)),  # Convert BatchEncoding instance to dictionary
    validation_label
)).batch(BATCH_SIZE).prefetch(1)
</code></pre>
<h3>Train</h3>
<pre><code># --------------------------------------------------------------------------------
# Train the model
# https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit
# Input data x can be a dict mapping input names to the corresponding array/tensors, 
# if the model has named inputs. Beware of the &quot;names&quot;. y should be consistent with x 
# (you cannot have Numpy inputs and tensor targets, or inversely). 
# --------------------------------------------------------------------------------
history = model.fit(
    x=X,    # dictionary 
    # y=Y,
    y=None,
    epochs=NUM_EPOCHS,
    batch_size=BATCH_SIZE,
    validation_data=V,
)
</code></pre>
<p>To implement the 1st approach, change the configuration as below.</p>
<pre><code>USE_CUSTOM_HEAD = False
</code></pre>
<p>Then <code>FREEZE_BASE</code> is changed to <code>False</code> and <code>LEARNING_RATE</code> is changed to <code>5e-5</code> which will run Further Pre-training on the base BERT model.</p>
<h3>Saving the model</h3>
<p>For the 3rd approach, saving the model will cause issues. The <a href="https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.save_pretrained" rel="noreferrer">save_pretrained</a> method of the Huggingface Model cannot be used as the model is not a direct sub class from of Huggingface <a href="https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel" rel="noreferrer">PreTrainedModel</a>.</p>
<p><a href="https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model" rel="noreferrer">Keras save_model</a> causes an error with the default <code>save_traces=True</code>, or causes a different error with <code>save_traces=True</code> when loading the model with <a href="https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model" rel="noreferrer">Keras load_model</a>.</p>
<pre><code>---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-71-01d66991d115&gt; in &lt;module&gt;()
----&gt; 1 tf.keras.models.load_model(MODEL_DIRECTORY)
 
11 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/saving/saved_model/load.py in _unable_to_call_layer_due_to_serialization_issue(layer, *unused_args, **unused_kwargs)
    865       'recorded when the object is called, and used when saving. To manually '
    866       'specify the input shape/dtype, decorate the call function with '
--&gt; 867       '`@tf.function(input_signature=...)`.'.format(layer.name, type(layer)))
    868 
    869 
 
ValueError: Cannot call custom layer tf_distil_bert_model of type &lt;class 'tensorflow.python.keras.saving.saved_model.load.TFDistilBertModel'&gt;, because the call function was not serialized to the SavedModel.Please try one of the following methods to fix this issue:
 
(1) Implement `get_config` and `from_config` in the layer/model class, and pass the object to the `custom_objects` argument when loading the model. For more details, see: https://www.tensorflow.org/guide/keras/save_and_serialize
 
(2) Ensure that the subclassed model or layer overwrites `call` and not `__call__`. The input shape and dtype will be automatically recorded when the object is called, and used when saving. To manually specify the input shape/dtype, decorate the call function with `@tf.function(input_signature=...)`.
</code></pre>
<p>Only <a href="https://www.tensorflow.org/api_docs/python/tf/keras/Model#save_weights" rel="noreferrer">Keras Model save_weights</a> worked as far as I tested.</p>
<h1>Experiments</h1>
<p>As far as I tested with <a href="https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge" rel="noreferrer">Toxic Comment Classification Challenge</a>, the 1st approach gave better recall (identify true toxic comment, true non-toxic comment). Code can be accessed as below. Please provide correction/suggestion if anything.</p>
<ul>
<li><a href="https://nbviewer.jupyter.org/github/omontasama/nlp-huggingface/blob/main/fine_tuning/huggingface_fine_tuning.ipynb" rel="noreferrer">Code for 1st and 3rd approach</a></li>
</ul>
<hr />
<h1>Related</h1>
<ul>
<li><a href="https://www.youtube.com/watch?v=_eSGWNqKeeY" rel="noreferrer">BERT Document Classification Tutorial with Code</a> - Fine tuning using TFDistilBertForSequenceClassification and Pytorch</li>
<li><a href="https://towardsdatascience.com/hugging-face-transformers-fine-tuning-distilbert-for-binary-classification-tasks-490f1d192379" rel="noreferrer">Hugging Face Transformers: Fine-tuning DistilBERT for Binary Classification Tasks</a> - Fine tuning using TFDistilBertModel</li>
</ul>
How to Fine-tune HuggingFace BERT model for Text Classification
<p>Is there a <em>Step by step explanation</em> on how to <strong>Fine-tune HuggingFace BERT</strong> model for text classification?</p>