18.3  Sentiment Analysis: Using Convolutional Neural Networks

In Chapter 6, we investigated mechanisms for processing two-dimensional image data with two-dimensional CNNs, which were applied to local features such as adjacent pixels. Though originally designed for computer vision, CNNs are also widely used for natural language processing. Simply put, just think of any text sequence as a one-dimensional image. In this way, one-dimensional CNNs can process local features such as \(n\)-grams in text.

In this section, we will use the textCNN model to demonstrate how to design a CNN architecture for representing single text (Kim 2014). Compared with Figure 18.2.1 that uses an RNN architecture with GloVe pretraining for sentiment analysis, the only difference in Figure 18.3.1 lies in the choice of the architecture.

Figure 18.3.1: This section feeds pretrained GloVe to a CNN-based architecture for sentiment analysis. 
from d2l import torch as d2l
import torch
from torch import nn

batch_size = 64
train_iter, test_iter, vocab = d2l.load_data_imdb(batch_size)
from d2l import tensorflow as d2l
import tensorflow as tf
import keras
import numpy as np

batch_size = 64
train_iter, test_iter, vocab = d2l.load_data_imdb(batch_size)
# d2l.load_array uses shuffle(buffer_size=1000), which is too small for
# the IMDb training set (25000 examples ordered as 12500 positives then
# 12500 negatives). Reshuffle the full dataset so each epoch sees a
# properly mixed class distribution, matching the PyTorch/JAX behavior.
train_iter = (train_iter.unbatch()
              .shuffle(25000, reshuffle_each_iteration=True)
              .batch(batch_size))
from d2l import jax as d2l
import jax
from jax import numpy as jnp
from flax import nnx
import optax
import numpy as np

batch_size = 256
train_iter, test_iter, vocab = d2l.load_data_imdb(batch_size)
from d2l import mxnet as d2l
from mxnet import gluon, init, np, npx
from mxnet.gluon import nn
npx.set_np()

batch_size = 64
train_iter, test_iter, vocab = d2l.load_data_imdb(batch_size)

18.3.1 One-Dimensional Convolutions

Before introducing the model, let’s see how a one-dimensional convolution works. Bear in mind that it is just a special case of a two-dimensional convolution based on the cross-correlation operation.

Figure 18.3.2: One-dimensional cross-correlation operation. The shaded portions are the first output element as well as the input and kernel tensor elements used for the output computation: \(0\times1+1\times2=2\).

As shown in Figure 18.3.2, in the one-dimensional case, the convolution window slides from left to right across the input tensor. During sliding, the input subtensor (e.g., \(0\) and \(1\) in Figure 18.3.2) contained in the convolution window at a certain position and the kernel tensor (e.g., \(1\) and \(2\) in Figure 18.3.2) are multiplied elementwise. The sum of these multiplications gives the single scalar value (e.g., \(0\times1+1\times2=2\) in Figure 18.3.2) at the corresponding position of the output tensor.

We implement one-dimensional cross-correlation in the following corr1d function. Given an input tensor X and a kernel tensor K, it returns the output tensor Y.

def corr1d(X, K):
    w = K.shape[0]
    Y = d2l.zeros((X.shape[0] - w + 1))
    for i in range(Y.shape[0]):
        Y[i] = (X[i: i + w] * K).sum()
    return Y
def corr1d(X, K):
    w = K.shape[0]
    Y = [tf.reduce_sum(X[i: i + w] * K) for i in range(X.shape[0] - w + 1)]
    return tf.stack(Y)
def corr1d(X, K):
    w = K.shape[0]
    Y = d2l.zeros((X.shape[0] - w + 1))
    for i in range(Y.shape[0]):
        Y = Y.at[i].set((X[i: i + w] * K).sum())
    return Y
def corr1d(X, K):
    w = K.shape[0]
    Y = d2l.zeros((X.shape[0] - w + 1))
    for i in range(Y.shape[0]):
        Y[i] = (X[i: i + w] * K).sum()
    return Y

We can construct the input tensor X and the kernel tensor K from Figure 18.3.2 to validate the output of the above one-dimensional cross-correlation implementation.

X, K = d2l.tensor([0, 1, 2, 3, 4, 5, 6]), d2l.tensor([1, 2])
corr1d(X, K)
tensor([ 2.,  5.,  8., 11., 14., 17.])
<tf.Tensor: shape=(6,), dtype=int32, numpy=array([ 2,  5,  8, 11, 14, 17], dtype=int32)>
Array([ 2.,  5.,  8., 11., 14., 17.], dtype=float32)
array([ 2.,  5.,  8., 11., 14., 17.])

For any one-dimensional input with multiple channels, the convolution kernel needs to have the same number of input channels. Then for each channel, perform a cross-correlation operation on the one-dimensional tensor of the input and the one-dimensional tensor of the convolution kernel, summing the results over all the channels to produce the one-dimensional output tensor. Figure 18.3.3 shows a one-dimensional cross-correlation operation with 3 input channels.

Figure 18.3.3: One-dimensional cross-correlation operation with 3 input channels. The shaded portions are the first output element as well as the input and kernel tensor elements used for the output computation: \(0\times1+1\times2+1\times3+2\times4+2\times(-1)+3\times(-3)=2\).

We can implement the one-dimensional cross-correlation operation for multiple input channels and validate the results in Figure 18.3.3.

def corr1d_multi_in(X, K):
    # First, iterate through the 0th dimension (channel dimension) of `X` and
    # `K`. Then, add them together
    return sum(corr1d(x, k) for x, k in zip(X, K))

X = d2l.tensor([[0, 1, 2, 3, 4, 5, 6],
              [1, 2, 3, 4, 5, 6, 7],
              [2, 3, 4, 5, 6, 7, 8]])
K = d2l.tensor([[1, 2], [3, 4], [-1, -3]])
corr1d_multi_in(X, K)
tensor([ 2.,  8., 14., 20., 26., 32.])
<tf.Tensor: shape=(6,), dtype=int32, numpy=array([ 2,  8, 14, 20, 26, 32], dtype=int32)>
Array([ 2.,  8., 14., 20., 26., 32.], dtype=float32)
array([ 2.,  8., 14., 20., 26., 32.])

Note that multi-input-channel one-dimensional cross-correlations are equivalent to single-input-channel two-dimensional cross-correlations. To illustrate, an equivalent form of the multi-input-channel one-dimensional cross-correlation in Figure 18.3.3 is the single-input-channel two-dimensional cross-correlation in Figure 18.3.4, where the height of the convolution kernel has to be the same as that of the input tensor.

Figure 18.3.4: Two-dimensional cross-correlation operation with a single input channel. The shaded portions are the first output element as well as the input and kernel tensor elements used for the output computation: \(2\times(-1)+3\times(-3)+1\times3+2\times4+0\times1+1\times2=2\).

Both the outputs in Figure 18.3.2 and Figure 18.3.3 have only one channel. Same as two-dimensional convolutions with multiple output channels described in Section 6.4.2, we can also specify multiple output channels for one-dimensional convolutions.

18.3.2 Max-Over-Time Pooling

Similarly, we can use pooling to extract the highest value from sequence representations as the most important feature across time steps. The max-over-time pooling used in textCNN works like the one-dimensional global max-pooling (Collobert et al. 2011). For a multi-channel input where each channel stores values at different time steps, the output at each channel is the maximum value for that channel. Note that the max-over-time pooling allows different numbers of time steps at different channels.

18.3.3 The textCNN Model

Using the one-dimensional convolution and max-over-time pooling, the textCNN model takes individual pretrained token representations as input, then obtains and transforms sequence representations for the downstream application.

For a single text sequence with \(n\) tokens represented by \(d\)-dimensional vectors, the width, height, and number of channels of the input tensor are \(n\), \(1\), and \(d\), respectively. The textCNN model transforms the input into the output as follows:

  1. Define multiple one-dimensional convolution kernels and perform convolution operations separately on the inputs. Convolution kernels with different widths may capture local features among different numbers of adjacent tokens.
  2. Perform max-over-time pooling on all the output channels, and then concatenate all the scalar pooling outputs as a vector.
  3. Transform the concatenated vector into the output categories using the fully connected layer. Dropout can be used for reducing overfitting.
Figure 18.3.5: The model architecture of textCNN.

Figure 18.3.5 illustrates the model architecture of textCNN with a concrete example. The input is a sentence with 11 tokens, where each token is represented by a 6-dimensional vector. So we have a 6-channel input with width 11. Define two one-dimensional convolution kernels of widths 2 and 4, with 4 and 5 output channels, respectively. They produce 4 output channels with width \(11-2+1=10\) and 5 output channels with width \(11-4+1=8\). Despite different widths of these 9 channels, the max-over-time pooling gives a concatenated 9-dimensional vector, which is finally transformed into a 2-dimensional output vector for binary sentiment predictions.

18.3.3.1 Defining the Model

We implement the textCNN model in the following class. Compared with the bidirectional RNN model in Section 18.2, besides replacing recurrent layers with convolutional layers, we also use two embedding layers: one with trainable weights and the other with fixed weights.

class TextCNN(nn.Module):
    def __init__(self, vocab_size, embed_size, kernel_sizes, num_channels,
                 **kwargs):
        super(TextCNN, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        # The embedding layer not to be trained
        self.constant_embedding = nn.Embedding(vocab_size, embed_size)
        self.dropout = nn.Dropout(0.5)
        self.decoder = nn.Linear(sum(num_channels), 2)
        # The max-over-time pooling layer has no parameters, so this instance
        # can be shared
        self.pool = nn.AdaptiveMaxPool1d(1)
        self.relu = nn.ReLU()
        # Create multiple one-dimensional convolutional layers
        self.convs = nn.ModuleList()
        for c, k in zip(num_channels, kernel_sizes):
            self.convs.append(nn.Conv1d(2 * embed_size, c, k))

    def forward(self, inputs):
        # Concatenate two embedding layer outputs with shape (batch size, no.
        # of tokens, token vector dimension) along vectors
        embeddings = torch.cat((
            self.embedding(inputs), self.constant_embedding(inputs)), dim=2)
        # Per the input format of one-dimensional convolutional layers,
        # rearrange the tensor so that the second dimension stores channels
        embeddings = embeddings.permute(0, 2, 1)
        # For each one-dimensional convolutional layer, after max-over-time
        # pooling, a tensor of shape (batch size, no. of channels, 1) is
        # obtained. Remove the last dimension and concatenate along channels
        encoding = torch.cat([
            torch.squeeze(self.relu(self.pool(conv(embeddings))), dim=-1)
            for conv in self.convs], dim=1)
        outputs = self.decoder(self.dropout(encoding))
        return outputs
class TextCNN(d2l.Classifier):
    def __init__(self, vocab_size, embed_size, kernel_sizes, num_channels,
                 **kwargs):
        super().__init__(**kwargs)
        self.embedding = keras.layers.Embedding(vocab_size, embed_size)
        # The embedding layer not to be trained
        self.constant_embedding = keras.layers.Embedding(vocab_size, embed_size)
        self.dropout = keras.layers.Dropout(0.5)
        self.decoder = keras.layers.Dense(2)
        # Create multiple one-dimensional convolutional layers
        self.convs = [keras.layers.Conv1D(c, k, activation='relu')
                      for c, k in zip(num_channels, kernel_sizes)]
        self.pool = keras.layers.GlobalMaxPooling1D()

    def call(self, inputs, training=False):
        # Concatenate two embedding layer outputs with shape
        # (batch_size, num_steps, 2 * embed_size) along the last axis
        embeddings = tf.concat(
            [self.embedding(inputs), self.constant_embedding(inputs)], axis=2)
        # For each convolutional layer, apply conv → global max pooling
        # and collect a (batch_size, num_channels) vector per kernel
        encoding = tf.concat(
            [self.pool(conv(embeddings)) for conv in self.convs], axis=1)
        outputs = self.decoder(self.dropout(encoding, training=training))
        return outputs
class TextCNN(nnx.Module):
    def __init__(self, vocab_size, embed_size, kernel_sizes, num_channels,
                 rngs=None):
        rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs
        self.embedding = nnx.Embed(vocab_size, embed_size, rngs=rngs)
        # The embedding layer not to be trained
        self.constant_embedding = nnx.Embed(vocab_size, embed_size, rngs=rngs)
        self.dropout = nnx.Dropout(0.5, rngs=rngs)
        self.decoder = nnx.Linear(sum(num_channels), 2, rngs=rngs)
        # Create multiple one-dimensional convolutional layers
        self.convs = nnx.List([
            nnx.Conv(2 * embed_size, c, kernel_size=(k,), rngs=rngs)
            for c, k in zip(num_channels, kernel_sizes)])

    def __call__(self, inputs):
        # Concatenate two embedding layer outputs with shape (batch size, no.
        # of tokens, token vector dimension) along vectors
        embeddings = jnp.concatenate((
            self.embedding(inputs), self.constant_embedding(inputs)), axis=2)
        # For Flax Conv, input shape is (batch, length, channels) which is
        # already the shape of embeddings
        # For each one-dimensional convolutional layer, after max-over-time
        # pooling, a tensor of shape (batch size, no. of channels) is obtained.
        # Concatenate along channels
        encoding = jnp.concatenate([
            jnp.max(nnx.relu(conv(embeddings)), axis=1)
            for conv in self.convs], axis=1)
        outputs = self.decoder(self.dropout(encoding))
        return outputs
class TextCNN(nn.Block):
    def __init__(self, vocab_size, embed_size, kernel_sizes, num_channels):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        # The embedding layer not to be trained
        self.constant_embedding = nn.Embedding(vocab_size, embed_size)
        self.dropout = nn.Dropout(0.5)
        self.decoder = nn.Dense(2)
        # The max-over-time pooling layer has no parameters, so this instance
        # can be shared
        self.pool = nn.GlobalMaxPool1D()
        # Create multiple one-dimensional convolutional layers
        self.convs = nn.Sequential()
        for c, k in zip(num_channels, kernel_sizes):
            self.convs.add(nn.Conv1D(c, k, activation='relu'))

    def forward(self, inputs):
        # Concatenate two embedding layer outputs with shape (batch size, no.
        # of tokens, token vector dimension) along vectors
        embeddings = np.concatenate((
            self.embedding(inputs), self.constant_embedding(inputs)), axis=2)
        # Per the input format of one-dimensional convolutional layers,
        # rearrange the tensor so that the second dimension stores channels
        embeddings = embeddings.transpose(0, 2, 1)
        # For each one-dimensional convolutional layer, after max-over-time
        # pooling, a tensor of shape (batch size, no. of channels, 1) is
        # obtained. Remove the last dimension and concatenate along channels
        encoding = np.concatenate([
            np.squeeze(self.pool(conv(embeddings)), axis=-1)
            for conv in self.convs], axis=1)
        outputs = self.decoder(self.dropout(encoding))
        return outputs

Let’s create a textCNN instance. It has 3 convolutional layers with kernel widths of 3, 4, and 5, all with 100 output channels.

embed_size, kernel_sizes, nums_channels = 100, [3, 4, 5], [100, 100, 100]
devices = d2l.try_all_gpus()
net = TextCNN(len(vocab), embed_size, kernel_sizes, nums_channels)

def init_weights(module):
    if type(module) in (nn.Linear, nn.Conv1d):
        nn.init.xavier_uniform_(module.weight)

net.apply(init_weights);
embed_size, kernel_sizes, nums_channels = 100, [3, 4, 5], [100, 100, 100]
devices = d2l.try_all_gpus()
net = TextCNN(len(vocab), embed_size, kernel_sizes, nums_channels)
# Build the model by calling it once on a dummy input
dummy_input = tf.zeros((1, 500), dtype=tf.int32)
net(dummy_input)
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[0.04802025, 0.01694743]], dtype=float32)>
embed_size, kernel_sizes, nums_channels = 100, [3, 4, 5], [100, 100, 100]
devices = d2l.try_all_gpus()
net = TextCNN(len(vocab), embed_size, kernel_sizes, nums_channels)
d2l.check_shape(nnx.view(net, deterministic=True)(
    jnp.ones((1, 500), dtype=jnp.int32)), (1, 2))
embed_size, kernel_sizes, nums_channels = 100, [3, 4, 5], [100, 100, 100]
devices = d2l.try_all_gpus()
net = TextCNN(len(vocab), embed_size, kernel_sizes, nums_channels)
net.initialize(init.Xavier(), ctx=devices)

18.3.3.2 Loading Pretrained Word Vectors

Same as Section 18.2, we load pretrained 100-dimensional GloVe embeddings as the initialized token representations. These token representations (embedding weights) will be trained in embedding and fixed in constant_embedding.

glove_embedding = d2l.TokenEmbedding('glove.6b.100d')
embeds = glove_embedding[vocab.idx_to_token]
net.embedding.weight.data.copy_(embeds)
net.constant_embedding.weight.data.copy_(embeds)
net.constant_embedding.weight.requires_grad = False
glove_embedding = d2l.TokenEmbedding('glove.6b.100d')
embeds = glove_embedding[vocab.idx_to_token]
net.embedding.set_weights([np.array(embeds)])
net.constant_embedding.set_weights([np.array(embeds)])
net.constant_embedding.trainable = False
glove_embedding = d2l.TokenEmbedding('glove.6b.100d')
embeds = glove_embedding[vocab.idx_to_token]
# Set pretrained embedding weights in the parameters
net.embedding.embedding[...] = jnp.array(embeds)
net.constant_embedding.embedding = nnx.data(jnp.array(embeds))
glove_embedding = d2l.TokenEmbedding('glove.6b.100d')
embeds = glove_embedding[vocab.idx_to_token]
net.embedding.weight.set_data(embeds)
net.constant_embedding.weight.set_data(embeds)
for p in net.constant_embedding.collect_params().values():
    p.grad_req = 'null'

18.3.3.3 Training and Evaluating the Model

Now we can train the textCNN model for sentiment analysis.

lr, num_epochs = 0.001, 5
trainer = torch.optim.Adam(net.parameters(), lr=lr)
loss = nn.CrossEntropyLoss(reduction="none")
d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)
loss 0.093, train acc 0.969, test acc 0.865
13063.9 examples/sec on [device(type='cuda', index=0)]

lr, num_epochs = 0.001, 5
net.compile(optimizer=keras.optimizers.Adam(lr),
            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
            metrics=['accuracy'])
net.fit(train_iter, validation_data=test_iter, epochs=num_epochs)
Epoch 1/5
      1/Unknown 10s 10s/step - accuracy: 0.5312 - loss: 2.2925
      9/Unknown 10s 6ms/step - accuracy: 0.5061 - loss: 1.6597
     18/Unknown 10s 6ms/step - accuracy: 0.5220 - loss: 1.4705
     27/Unknown 10s 6ms/step - accuracy: 0.5281 - loss: 1.3641
     36/Unknown 10s 6ms/step - accuracy: 0.5311 - loss: 1.2932
     45/Unknown 10s 6ms/step - accuracy: 0.5344 - loss: 1.2395
     53/Unknown 10s 6ms/step - accuracy: 0.5378 - loss: 1.2004
     61/Unknown 10s 6ms/step - accuracy: 0.5421 - loss: 1.1662
     70/Unknown 11s 6ms/step - accuracy: 0.5469 - loss: 1.1327
     79/Unknown 11s 6ms/step - accuracy: 0.5519 - loss: 1.1032
     88/Unknown 11s 6ms/step - accuracy: 0.5564 - loss: 1.0775
     97/Unknown 11s 6ms/step - accuracy: 0.5606 - loss: 1.0546
    106/Unknown 11s 6ms/step - accuracy: 0.5647 - loss: 1.0338
    115/Unknown 11s 6ms/step - accuracy: 0.5689 - loss: 1.0149
    124/Unknown 11s 6ms/step - accuracy: 0.5729 - loss: 0.9975
    133/Unknown 11s 6ms/step - accuracy: 0.5770 - loss: 0.9814
    142/Unknown 11s 6ms/step - accuracy: 0.5810 - loss: 0.9664
    152/Unknown 11s 6ms/step - accuracy: 0.5853 - loss: 0.9509
    162/Unknown 11s 6ms/step - accuracy: 0.5895 - loss: 0.9364
    172/Unknown 11s 6ms/step - accuracy: 0.5935 - loss: 0.9228
    181/Unknown 11s 6ms/step - accuracy: 0.5971 - loss: 0.9113
    190/Unknown 11s 6ms/step - accuracy: 0.6004 - loss: 0.9004
    199/Unknown 11s 6ms/step - accuracy: 0.6037 - loss: 0.8902
    208/Unknown 11s 6ms/step - accuracy: 0.6068 - loss: 0.8804
    217/Unknown 11s 6ms/step - accuracy: 0.6099 - loss: 0.8713
    225/Unknown 11s 6ms/step - accuracy: 0.6125 - loss: 0.8635
    233/Unknown 11s 6ms/step - accuracy: 0.6150 - loss: 0.8560
    242/Unknown 12s 6ms/step - accuracy: 0.6177 - loss: 0.8481
    251/Unknown 12s 6ms/step - accuracy: 0.6204 - loss: 0.8405
    260/Unknown 12s 6ms/step - accuracy: 0.6230 - loss: 0.8332
    269/Unknown 12s 6ms/step - accuracy: 0.6255 - loss: 0.8262
    277/Unknown 12s 6ms/step - accuracy: 0.6277 - loss: 0.8203
    286/Unknown 12s 6ms/step - accuracy: 0.6300 - loss: 0.8138
    293/Unknown 12s 6ms/step - accuracy: 0.6318 - loss: 0.8090
    301/Unknown 12s 6ms/step - accuracy: 0.6338 - loss: 0.8036
    310/Unknown 12s 6ms/step - accuracy: 0.6360 - loss: 0.7978
    319/Unknown 12s 6ms/step - accuracy: 0.6381 - loss: 0.7923
    329/Unknown 12s 6ms/step - accuracy: 0.6404 - loss: 0.7863
    338/Unknown 12s 6ms/step - accuracy: 0.6424 - loss: 0.7811
    347/Unknown 12s 6ms/step - accuracy: 0.6443 - loss: 0.7760
    357/Unknown 12s 6ms/step - accuracy: 0.6464 - loss: 0.7707
    366/Unknown 12s 6ms/step - accuracy: 0.6483 - loss: 0.7660
    374/Unknown 12s 6ms/step - accuracy: 0.6499 - loss: 0.7619
    383/Unknown 12s 6ms/step - accuracy: 0.6516 - loss: 0.7575
    391/Unknown 19s 23ms/step - accuracy: 0.6532 - loss: 0.7537
/home/smola/d2l-neu/.venv-tensorflow/lib/python3.12/site-packages/keras/src/trainers/epoch_iterator.py:164: UserWarning: Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches. You may need to use the `.repeat()` function when building your dataset.
  self._interrupted_warning()
391/391 ━━━━━━━━━━━━━━━━━━━━ 23s 34ms/step - accuracy: 0.7268 - loss: 0.5696 - val_accuracy: 0.8362 - val_loss: 0.3704
Epoch 2/5
  1/391 ━━━━━━━━━━━━━━━━━━━━ 40s 104ms/step - accuracy: 0.8750 - loss: 0.3433
 10/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8460 - loss: 0.3675   
 19/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8458 - loss: 0.3690
 28/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8411 - loss: 0.3750
 37/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8392 - loss: 0.3775
 46/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8387 - loss: 0.3780
 56/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8389 - loss: 0.3767
 65/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8390 - loss: 0.3755
 75/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8389 - loss: 0.3744
 84/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8389 - loss: 0.3735
 93/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8391 - loss: 0.3725
102/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8393 - loss: 0.3715
111/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8394 - loss: 0.3707
120/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8395 - loss: 0.3701
129/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8397 - loss: 0.3694
138/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8398 - loss: 0.3689
148/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8398 - loss: 0.3685
158/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8397 - loss: 0.3683
168/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8396 - loss: 0.3681
178/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8395 - loss: 0.3679
188/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8394 - loss: 0.3678
198/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8393 - loss: 0.3678
208/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8393 - loss: 0.3677
218/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3676
228/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3676
237/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8392 - loss: 0.3675
246/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3674
255/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3673
265/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3672
275/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3671
285/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3671
295/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8392 - loss: 0.3670
305/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8392 - loss: 0.3669
314/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8392 - loss: 0.3668
323/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8392 - loss: 0.3667
332/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3665
341/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8393 - loss: 0.3664
350/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8394 - loss: 0.3662
360/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8394 - loss: 0.3661
369/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8394 - loss: 0.3660
379/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8395 - loss: 0.3658
388/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8396 - loss: 0.3657
391/391 ━━━━━━━━━━━━━━━━━━━━ 3s 8ms/step - accuracy: 0.8421 - loss: 0.3604 - val_accuracy: 0.8592 - val_loss: 0.3274
Epoch 3/5
  1/391 ━━━━━━━━━━━━━━━━━━━━ 39s 100ms/step - accuracy: 0.8438 - loss: 0.3754
 10/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8831 - loss: 0.2984   
 19/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8807 - loss: 0.2954
 27/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8815 - loss: 0.2906
 36/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.8835 - loss: 0.2847
 46/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8849 - loss: 0.2811
 55/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8861 - loss: 0.2786
 64/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8868 - loss: 0.2767
 73/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8874 - loss: 0.2752
 83/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8881 - loss: 0.2736
 92/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8886 - loss: 0.2723
102/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8893 - loss: 0.2708
112/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8899 - loss: 0.2694
122/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8906 - loss: 0.2680
132/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8912 - loss: 0.2667
141/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8917 - loss: 0.2658
150/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8921 - loss: 0.2649
159/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8924 - loss: 0.2643
168/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8927 - loss: 0.2637
177/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8929 - loss: 0.2632
187/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8931 - loss: 0.2628
197/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8933 - loss: 0.2623
206/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.8934 - loss: 0.2620
215/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8936 - loss: 0.2617
224/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8937 - loss: 0.2615
233/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8938 - loss: 0.2612
242/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8938 - loss: 0.2610
251/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8939 - loss: 0.2609
260/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8940 - loss: 0.2607
269/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8941 - loss: 0.2606
280/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8941 - loss: 0.2604
290/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8942 - loss: 0.2603
299/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8942 - loss: 0.2602
308/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8942 - loss: 0.2602
317/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8942 - loss: 0.2602
326/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8942 - loss: 0.2602
336/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8943 - loss: 0.2601
345/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8943 - loss: 0.2601
354/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8943 - loss: 0.2601
363/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8943 - loss: 0.2601
373/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8943 - loss: 0.2601
382/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.8943 - loss: 0.2601
391/391 ━━━━━━━━━━━━━━━━━━━━ 3s 8ms/step - accuracy: 0.8946 - loss: 0.2589 - val_accuracy: 0.8710 - val_loss: 0.3042
Epoch 4/5
  1/391 ━━━━━━━━━━━━━━━━━━━━ 1:13 189ms/step - accuracy: 0.9844 - loss: 0.1185
 11/391 ━━━━━━━━━━━━━━━━━━━━ 2s 5ms/step - accuracy: 0.9357 - loss: 0.1780    
 20/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9313 - loss: 0.1821
 29/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9315 - loss: 0.1807
 38/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9324 - loss: 0.1800
 47/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9329 - loss: 0.1796
 57/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9332 - loss: 0.1791
 65/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9339 - loss: 0.1782
 73/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9346 - loss: 0.1772
 82/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9352 - loss: 0.1761
 90/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9355 - loss: 0.1756
 99/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9360 - loss: 0.1748
108/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9364 - loss: 0.1739
117/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9369 - loss: 0.1731
126/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9373 - loss: 0.1722
135/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9376 - loss: 0.1715
144/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9380 - loss: 0.1709
153/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9383 - loss: 0.1703
163/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9385 - loss: 0.1698
172/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9387 - loss: 0.1695
180/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9388 - loss: 0.1693
189/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9389 - loss: 0.1690
198/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9391 - loss: 0.1686
207/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9392 - loss: 0.1683
216/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9393 - loss: 0.1680
225/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9394 - loss: 0.1677
234/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9396 - loss: 0.1674
243/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9397 - loss: 0.1672
251/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9398 - loss: 0.1669
259/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9399 - loss: 0.1667
267/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9400 - loss: 0.1664
275/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9401 - loss: 0.1662
284/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9402 - loss: 0.1660
293/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9403 - loss: 0.1658
302/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9403 - loss: 0.1656
311/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9404 - loss: 0.1654
320/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9404 - loss: 0.1653
329/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9405 - loss: 0.1651
338/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9405 - loss: 0.1650
347/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9405 - loss: 0.1649
356/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9405 - loss: 0.1648
364/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9406 - loss: 0.1647
372/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9406 - loss: 0.1646
380/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9406 - loss: 0.1645
389/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9406 - loss: 0.1645
391/391 ━━━━━━━━━━━━━━━━━━━━ 4s 9ms/step - accuracy: 0.9405 - loss: 0.1615 - val_accuracy: 0.8628 - val_loss: 0.3367
Epoch 5/5
  1/391 ━━━━━━━━━━━━━━━━━━━━ 1:24 216ms/step - accuracy: 0.9531 - loss: 0.0929
 10/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9683 - loss: 0.0901    
 18/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9709 - loss: 0.0874
 27/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9705 - loss: 0.0895
 36/391 ━━━━━━━━━━━━━━━━━━━━ 2s 6ms/step - accuracy: 0.9707 - loss: 0.0903
 47/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9710 - loss: 0.0902
 57/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9713 - loss: 0.0902
 67/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9710 - loss: 0.0910
 76/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9707 - loss: 0.0916
 85/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9705 - loss: 0.0921
 94/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9705 - loss: 0.0923
103/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9706 - loss: 0.0924
112/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9706 - loss: 0.0925
121/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9706 - loss: 0.0925
130/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9707 - loss: 0.0925
139/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9707 - loss: 0.0925
148/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9708 - loss: 0.0924
157/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9709 - loss: 0.0922
166/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9710 - loss: 0.0920
175/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9711 - loss: 0.0918
185/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9712 - loss: 0.0916
195/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9713 - loss: 0.0915
205/391 ━━━━━━━━━━━━━━━━━━━━ 1s 6ms/step - accuracy: 0.9714 - loss: 0.0913
215/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9714 - loss: 0.0912
225/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9715 - loss: 0.0911
235/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9715 - loss: 0.0911
245/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9715 - loss: 0.0910
254/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9715 - loss: 0.0910
264/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9714 - loss: 0.0910
274/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9714 - loss: 0.0910
283/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9714 - loss: 0.0911
293/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9714 - loss: 0.0911
302/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9713 - loss: 0.0911
312/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9713 - loss: 0.0912
322/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9713 - loss: 0.0912
332/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9712 - loss: 0.0912
341/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9712 - loss: 0.0912
350/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9712 - loss: 0.0912
361/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9712 - loss: 0.0912
371/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9711 - loss: 0.0912
381/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9711 - loss: 0.0912
391/391 ━━━━━━━━━━━━━━━━━━━━ 0s 6ms/step - accuracy: 0.9710 - loss: 0.0912
391/391 ━━━━━━━━━━━━━━━━━━━━ 3s 8ms/step - accuracy: 0.9693 - loss: 0.0917 - val_accuracy: 0.8626 - val_loss: 0.3731
<keras.src.callbacks.history.History at 0x7ab2e86957f0>
lr, num_epochs = 0.001, 5
optimizer = nnx.Optimizer(net, optax.adam(lr), wrt=nnx.Param)
loss_fn = optax.softmax_cross_entropy_with_integer_labels

@nnx.jit
def train_step(net, optimizer, X, y):
    def compute_loss(model):
        logits = model(X)
        return loss_fn(logits, y).mean(), logits
    (loss, logits), grads = nnx.value_and_grad(
        compute_loss, has_aux=True)(net)
    optimizer.update(net, grads)
    return loss, logits

for epoch in range(num_epochs):
    loss_sum, train_correct, num_train = (
        jnp.array(0.0), jnp.array(0), 0)
    for X, y in train_iter:
        l, logits = train_step(net, optimizer, X, y)
        loss_sum += l * len(y)
        train_correct += (logits.argmax(axis=-1) == y).sum()
        num_train += len(y)
    # Evaluate
    correct, total = jnp.array(0), 0
    for X, y in test_iter:
        logits = nnx.view(net, deterministic=True)(X)
        correct += (logits.argmax(axis=-1) == y).sum()
        total += len(y)
    loss_sum, train_correct, correct = (
        float(loss_sum), int(train_correct), int(correct))
    print(f'epoch {epoch + 1}, loss {loss_sum / num_train:.3f}, '
          f'train acc {train_correct / num_train:.3f}, '
          f'test acc {correct / total:.3f}')
epoch 1, loss 0.644, train acc 0.664, test acc 0.812
epoch 2, loss 0.413, train acc 0.811, test acc 0.842
epoch 3, loss 0.339, train acc 0.851, test acc 0.856
epoch 4, loss 0.280, train acc 0.886, test acc 0.866
epoch 5, loss 0.219, train acc 0.916, test acc 0.875
lr, num_epochs = 0.001, 5
trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': lr})
loss = gluon.loss.SoftmaxCrossEntropyLoss()
d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)
loss 0.094, train acc 0.967, test acc 0.865
1097.4 examples/sec on [gpu(0)]

Below we use the trained model to predict the sentiment for two simple sentences.

d2l.predict_sentiment(net, vocab, 'this movie is so great')
'positive'
d2l.predict_sentiment(net, vocab, 'this movie is so great')
'positive'
tokens = jnp.array(vocab['this movie is so great'.split()])
logits = nnx.view(net, deterministic=True)(tokens.reshape(1, -1))
'positive' if int(jnp.argmax(logits, axis=1)[0]) == 1 else 'negative'
'positive'
d2l.predict_sentiment(net, vocab, 'this movie is so great')
'positive'
d2l.predict_sentiment(net, vocab, 'this movie is so bad')
'negative'
d2l.predict_sentiment(net, vocab, 'this movie is so bad')
'negative'
tokens = jnp.array(vocab['this movie is so bad'.split()])
logits = nnx.view(net, deterministic=True)(tokens.reshape(1, -1))
'positive' if int(jnp.argmax(logits, axis=1)[0]) == 1 else 'negative'
'negative'
d2l.predict_sentiment(net, vocab, 'this movie is so bad')
'negative'

18.3.4 Summary

  • One-dimensional CNNs can process local features such as \(n\)-grams in text.
  • Multi-input-channel one-dimensional cross-correlations are equivalent to single-input-channel two-dimensional cross-correlations.
  • The max-over-time pooling allows different numbers of time steps at different channels.
  • The textCNN model transforms individual token representations into downstream application outputs using one-dimensional convolutional layers and max-over-time pooling layers.

18.3.5 Exercises

  1. Tune hyperparameters and compare the two architectures for sentiment analysis in Section 18.2 and in this section, such as in classification accuracy and computational efficiency.
  2. Can you further improve the classification accuracy of the model by using the methods introduced in the exercises of Section 18.2?
  3. Add positional encoding in the input representations. Does it improve the classification accuracy?