Text classification with Switch Transformer

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/nlp/text_classification_with_switch_transformer/" />  <meta property="og:title" content="Keras documentation: Text classification with Switch Transformer"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Text classification with Switch Transformer"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Text classification with Switch Transformer</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink active" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink2" href="/examples/nlp/text_classification_from_scratch/">Text classification from scratch</a> <a class="nav-sublink2" href="/examples/nlp/active_learning_review_classification/">Review Classification using Active Learning</a> <a class="nav-sublink2" href="/examples/nlp/fnet_classification_with_keras_hub/">Text Classification using FNet</a> <a class="nav-sublink2" href="/examples/nlp/multi_label_classification/">Large-scale multi-label text classification</a> <a class="nav-sublink2" href="/examples/nlp/text_classification_with_transformer/">Text classification with Transformer</a> <a class="nav-sublink2 active" href="/examples/nlp/text_classification_with_switch_transformer/">Text classification with Switch Transformer</a> <a class="nav-sublink2" href="/examples/nlp/tweet-classification-using-tfdf/">Text classification using Decision Forests and pretrained embeddings</a> <a class="nav-sublink2" href="/examples/nlp/pretrained_word_embeddings/">Using pre-trained word embeddings</a> <a class="nav-sublink2" href="/examples/nlp/bidirectional_lstm_imdb/">Bidirectional LSTM on IMDB</a> <a class="nav-sublink2" href="/examples/nlp/data_parallel_training_with_keras_hub/">Data Parallel Training with KerasHub and tf.distribute</a> <a class="nav-sublink2" href="/examples/nlp/neural_machine_translation_with_keras_hub/">English-to-Spanish translation with KerasHub</a> <a class="nav-sublink2" href="/examples/nlp/neural_machine_translation_with_transformer/">English-to-Spanish translation with a sequence-to-sequence Transformer</a> <a class="nav-sublink2" href="/examples/nlp/lstm_seq2seq/">Character-level recurrent sequence-to-sequence model</a> <a class="nav-sublink2" href="/examples/nlp/multimodal_entailment/">Multimodal entailment</a> <a class="nav-sublink2" href="/examples/nlp/ner_transformers/">Named Entity Recognition using Transformers</a> <a class="nav-sublink2" href="/examples/nlp/text_extraction_with_bert/">Text Extraction with BERT</a> <a class="nav-sublink2" href="/examples/nlp/addition_rnn/">Sequence to sequence learning for performing number addition</a> <a class="nav-sublink2" href="/examples/nlp/semantic_similarity_with_keras_hub/">Semantic Similarity with KerasHub</a> <a class="nav-sublink2" href="/examples/nlp/semantic_similarity_with_bert/">Semantic Similarity with BERT</a> <a class="nav-sublink2" href="/examples/nlp/sentence_embeddings_with_sbert/">Sentence embeddings using Siamese RoBERTa-networks</a> <a class="nav-sublink2" href="/examples/nlp/masked_language_modeling/">End-to-end Masked Language Modeling with BERT</a> <a class="nav-sublink2" href="/examples/nlp/abstractive_summarization_with_bart/">Abstractive Text Summarization with BART</a> <a class="nav-sublink2" href="/examples/nlp/pretraining_BERT/">Pretraining BERT with Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora/">Parameter-efficient fine-tuning of GPT-2 with LoRA</a> <a class="nav-sublink2" href="/examples/nlp/mlm_training_tpus/">Training a language model from scratch with 🤗 Transformers and TPUs</a> <a class="nav-sublink2" href="/examples/nlp/multiple_choice_task_with_transfer_learning/">MultipleChoice Task with Transfer Learning</a> <a class="nav-sublink2" href="/examples/nlp/question_answering/">Question Answering with Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/nlp/t5_hf_summarization/">Abstractive Summarization with Hugging Face Transformers</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparameter Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> <a class="nav-link" href="/keras_cv/" role="tab" aria-selected="">KerasCV: Computer Vision Workflows</a> <a class="nav-link" href="/keras_nlp/" role="tab" aria-selected="">KerasNLP: Natural Language Workflows</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/nlp/'>Natural Language Processing</a> / Text classification with Switch Transformer </div> <div class='k-content'> <h1 id="text-classification-with-switch-transformer">Text classification with Switch Transformer</h1> Author: <a href="https://www.linkedin.com/in/khalid-salama-24403144/">Khalid Salama</a> Date created: 2020/05/10 Last modified: 2021/02/15 Description: Implement a Switch Transformer for text classification. <div class='example_version_banner keras_3'>ⓘ This example uses Keras 3</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/nlp/ipynb/text_classification_with_switch_transformer.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/nlp/text_classification_with_switch_transformer.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> This example demonstrates the implementation of the <a href="https://arxiv.org/abs/2101.03961">Switch Transformer</a> model for text classification. The Switch Transformer replaces the feedforward network (FFN) layer in the standard Transformer with a Mixture of Expert (MoE) routing layer, where each expert operates independently on the tokens in the sequence. This allows increasing the model size without increasing the computation needed to process each example. Note that, for training the Switch Transformer efficiently, data and model parallelism need to be applied, so that expert modules can run simultaneously, each on its own accelerator. While the implementation described in the paper uses the <a href="https://github.com/tensorflow/mesh">TensorFlow Mesh</a> framework for distributed training, this example presents a simple, non-distributed implementation of the Switch Transformer model for demonstration purposes. <hr /> <h2 id="setup">Setup</h2> <div class="codehilite"><pre><code>import keras from keras import ops from keras import layers </code></pre></div> <hr /> <h2 id="download-and-prepare-dataset">Download and prepare dataset</h2> <div class="codehilite"><pre><code>vocab_size = 20000 # Only consider the top 20k words num_tokens_per_example = 200 # Only consider the first 200 words of each movie review (x_train, y_train), (x_val, y_val) = keras.datasets.imdb.load_data(num_words=vocab_size) print(len(x_train), "Training sequences") print(len(x_val), "Validation sequences") x_train = keras.utils.pad_sequences(x_train, maxlen=num_tokens_per_example) x_val = keras.utils.pad_sequences(x_val, maxlen=num_tokens_per_example) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>25000 Training sequences 25000 Validation sequences </code></pre></div> </div> <hr /> <h2 id="define-hyperparameters">Define hyperparameters</h2> <div class="codehilite"><pre><code>embed_dim = 32 # Embedding size for each token. num_heads = 2 # Number of attention heads ff_dim = 32 # Hidden layer size in feedforward network. num_experts = 10 # Number of experts used in the Switch Transformer. batch_size = 50 # Batch size. learning_rate = 0.001 # Learning rate. dropout_rate = 0.25 # Dropout rate. num_epochs = 3 # Number of epochs. num_tokens_per_batch = ( batch_size * num_tokens_per_example ) # Total number of tokens per batch. print(f"Number of tokens per batch: {num_tokens_per_batch}") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Number of tokens per batch: 10000 </code></pre></div> </div> <hr /> <h2 id="implement-token-amp-position-embedding-layer">Implement token & position embedding layer</h2> It consists of two separate embedding layers, one for tokens, one for token index (positions). <div class="codehilite"><pre><code>class TokenAndPositionEmbedding(layers.Layer): def __init__(self, maxlen, vocab_size, embed_dim): super().__init__() self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim) self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim) def call(self, x): maxlen = ops.shape(x)[-1] positions = ops.arange(start=0, stop=maxlen, step=1) positions = self.pos_emb(positions) x = self.token_emb(x) return x + positions </code></pre></div> <hr /> <h2 id="implement-the-feedforward-network">Implement the feedforward network</h2> This is used as the Mixture of Experts in the Switch Transformer. <div class="codehilite"><pre><code>def create_feedforward_network(ff_dim, embed_dim, name=None): return keras.Sequential( [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim)], name=name ) </code></pre></div> <hr /> <h2 id="implement-the-loadbalanced-loss">Implement the load-balanced loss</h2> This is an auxiliary loss to encourage a balanced load across experts. <div class="codehilite"><pre><code>def load_balanced_loss(router_probs, expert_mask): # router_probs [tokens_per_batch, num_experts] is the probability assigned for # each expert per token. expert_mask [tokens_per_batch, num_experts] contains # the expert with the highest router probability in one−hot format. num_experts = ops.shape(expert_mask)[-1] # Get the fraction of tokens routed to each expert. # density is a vector of length num experts that sums to 1. density = ops.mean(expert_mask, axis=0) # Get fraction of probability mass assigned to each expert from the router # across all tokens. density_proxy is a vector of length num experts that sums to 1. density_proxy = ops.mean(router_probs, axis=0) # Want both vectors to have uniform allocation (1/num experts) across all # num_expert elements. The two vectors will be pushed towards uniform allocation # when the dot product is minimized. loss = ops.mean(density_proxy * density) * ops.cast((num_experts**2), "float32") return loss </code></pre></div> <h3 id="implement-the-router-as-a-layer">Implement the router as a layer</h3> <div class="codehilite"><pre><code>class Router(layers.Layer): def __init__(self, num_experts, expert_capacity): self.num_experts = num_experts self.route = layers.Dense(units=num_experts) self.expert_capacity = expert_capacity super().__init__() def call(self, inputs, training=False): # inputs shape: [tokens_per_batch, embed_dim] # router_logits shape: [tokens_per_batch, num_experts] router_logits = self.route(inputs) if training: # Add noise for exploration across experts. router_logits += keras.random.uniform( shape=router_logits.shape, minval=0.9, maxval=1.1 ) # Probabilities for each token of what expert it should be sent to. router_probs = keras.activations.softmax(router_logits, axis=-1) # Get the top−1 expert for each token. expert_gate is the top−1 probability # from the router for each token. expert_index is what expert each token # is going to be routed to. expert_gate, expert_index = ops.top_k(router_probs, k=1) # expert_mask shape: [tokens_per_batch, num_experts] expert_mask = ops.one_hot(expert_index, self.num_experts) # Compute load balancing loss. aux_loss = load_balanced_loss(router_probs, expert_mask) self.add_loss(aux_loss) # Experts have a fixed capacity, ensure we do not exceed it. Construct # the batch indices, to each expert, with position in expert make sure that # not more that expert capacity examples can be routed to each expert. position_in_expert = ops.cast( ops.cumsum(expert_mask, axis=0) * expert_mask, "int32" ) # Keep only tokens that fit within expert capacity. expert_mask *= ops.cast( ops.less(ops.cast(position_in_expert, "int32"), self.expert_capacity), "float32", ) expert_mask_flat = ops.sum(expert_mask, axis=-1) # Mask out the experts that have overflowed the expert capacity. expert_gate *= expert_mask_flat # Combine expert outputs and scaling with router probability. # combine_tensor shape: [tokens_per_batch, num_experts, expert_capacity] combined_tensor = ops.expand_dims( expert_gate * expert_mask_flat * ops.squeeze(ops.one_hot(expert_index, self.num_experts), 1), -1, ) * ops.squeeze(ops.one_hot(position_in_expert, self.expert_capacity), 1) # Create binary dispatch_tensor [tokens_per_batch, num_experts, expert_capacity] # that is 1 if the token gets routed to the corresponding expert. dispatch_tensor = ops.cast(combined_tensor, "float32") return dispatch_tensor, combined_tensor </code></pre></div> <h3 id="implement-a-switch-layer">Implement a Switch layer</h3> <div class="codehilite"><pre><code>class Switch(layers.Layer): def __init__( self, num_experts, embed_dim, ff_dim, num_tokens_per_batch, capacity_factor=1 ): self.num_experts = num_experts self.embed_dim = embed_dim self.experts = [ create_feedforward_network(ff_dim, embed_dim) for _ in range(num_experts) ] self.expert_capacity = num_tokens_per_batch // self.num_experts self.router = Router(self.num_experts, self.expert_capacity) super().__init__() def call(self, inputs): batch_size = ops.shape(inputs)[0] num_tokens_per_example = ops.shape(inputs)[1] # inputs shape: [num_tokens_per_batch, embed_dim] inputs = ops.reshape(inputs, [num_tokens_per_batch, self.embed_dim]) # dispatch_tensor shape: [expert_capacity, num_experts, tokens_per_batch] # combine_tensor shape: [tokens_per_batch, num_experts, expert_capacity] dispatch_tensor, combine_tensor = self.router(inputs) # expert_inputs shape: [num_experts, expert_capacity, embed_dim] expert_inputs = ops.einsum("ab,acd->cdb", inputs, dispatch_tensor) expert_inputs = ops.reshape( expert_inputs, [self.num_experts, self.expert_capacity, self.embed_dim] ) # Dispatch to experts expert_input_list = ops.unstack(expert_inputs, axis=0) expert_output_list = [ self.experts[idx](expert_input) for idx, expert_input in enumerate(expert_input_list) ] # expert_outputs shape: [expert_capacity, num_experts, embed_dim] expert_outputs = ops.stack(expert_output_list, axis=1) # expert_outputs_combined shape: [tokens_per_batch, embed_dim] expert_outputs_combined = ops.einsum( "abc,xba->xc", expert_outputs, combine_tensor ) # output shape: [batch_size, num_tokens_per_example, embed_dim] outputs = ops.reshape( expert_outputs_combined, [batch_size, num_tokens_per_example, self.embed_dim], ) return outputs </code></pre></div> <hr /> <h2 id="implement-a-transformer-block-layer">Implement a Transformer block layer</h2> <div class="codehilite"><pre><code>class TransformerBlock(layers.Layer): def __init__(self, embed_dim, num_heads, ffn, dropout_rate=0.1): super().__init__() self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim) # The ffn can be either a standard feedforward network or a switch # layer with a Mixture of Experts. self.ffn = ffn self.layernorm1 = layers.LayerNormalization(epsilon=1e-6) self.layernorm2 = layers.LayerNormalization(epsilon=1e-6) self.dropout1 = layers.Dropout(dropout_rate) self.dropout2 = layers.Dropout(dropout_rate) def call(self, inputs, training=False): attn_output = self.att(inputs, inputs) attn_output = self.dropout1(attn_output, training=training) out1 = self.layernorm1(inputs + attn_output) ffn_output = self.ffn(out1) ffn_output = self.dropout2(ffn_output, training=training) return self.layernorm2(out1 + ffn_output) </code></pre></div> <hr /> <h2 id="implement-the-classifier">Implement the classifier</h2> The <code>TransformerBlock</code> layer outputs one vector for each time step of our input sequence. Here, we take the mean across all time steps and use a feedforward network on top of it to classify text. <div class="codehilite"><pre><code>def create_classifier(): switch = Switch(num_experts, embed_dim, ff_dim, num_tokens_per_batch) transformer_block = TransformerBlock(embed_dim // num_heads, num_heads, switch) inputs = layers.Input(shape=(num_tokens_per_example,)) embedding_layer = TokenAndPositionEmbedding( num_tokens_per_example, vocab_size, embed_dim ) x = embedding_layer(inputs) x = transformer_block(x) x = layers.GlobalAveragePooling1D()(x) x = layers.Dropout(dropout_rate)(x) x = layers.Dense(ff_dim, activation="relu")(x) x = layers.Dropout(dropout_rate)(x) outputs = layers.Dense(2, activation="softmax")(x) classifier = keras.Model(inputs=inputs, outputs=outputs) return classifier </code></pre></div> <hr /> <h2 id="train-and-evaluate-the-model">Train and evaluate the model</h2> <div class="codehilite"><pre><code>def run_experiment(classifier): classifier.compile( optimizer=keras.optimizers.Adam(learning_rate), loss="sparse_categorical_crossentropy", metrics=["accuracy"], ) history = classifier.fit( x_train, y_train, batch_size=batch_size, epochs=num_epochs, validation_data=(x_val, y_val), ) return history classifier = create_classifier() run_experiment(classifier) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/3 500/500 ━━━━━━━━━━━━━━━━━━━━ 251s 485ms/step - accuracy: 0.7121 - loss: 1.5394 - val_accuracy: 0.8748 - val_loss: 1.2891 Epoch 2/3 500/500 ━━━━━━━━━━━━━━━━━━━━ 240s 480ms/step - accuracy: 0.9243 - loss: 1.2063 - val_accuracy: 0.8752 - val_loss: 1.3090 Epoch 3/3 500/500 ━━━━━━━━━━━━━━━━━━━━ 242s 485ms/step - accuracy: 0.9572 - loss: 1.1222 - val_accuracy: 0.8614 - val_loss: 1.3744 <keras.src.callbacks.history.History at 0x7efb79d82a90> </code></pre></div> </div> <hr /> <h2 id="conclusion">Conclusion</h2> Compared to the standard Transformer architecture, the Switch Transformer can have a much larger number of parameters, leading to increased model capacity, while maintaining a reasonable computational cost. </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#text-classification-with-switch-transformer'>Text classification with Switch Transformer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#download-and-prepare-dataset'>Download and prepare dataset</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#define-hyperparameters'>Define hyperparameters</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#implement-token-amp-position-embedding-layer'>Implement token & position embedding layer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#implement-the-feedforward-network'>Implement the feedforward network</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#implement-the-loadbalanced-loss'>Implement the load-balanced loss</a> </div> <div class='k-outline-depth-3'> <a href='#implement-the-router-as-a-layer'>Implement the router as a layer</a> </div> <div class='k-outline-depth-3'> <a href='#implement-a-switch-layer'>Implement a Switch layer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#implement-a-transformer-block-layer'>Implement a Transformer block layer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#implement-the-classifier'>Implement the classifier</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#train-and-evaluate-the-model'>Train and evaluate the model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#conclusion'>Conclusion</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>

CINXE.COM

Text classification with Switch Transformer