Apache MXNet - Gluon

Apache MXNet - Gluon API

The Gluon API (mxnet.gluon) is a high-level interface in Apache MXNet designed for ease of use, flexibility, and rapid prototyping of deep learning models. It provides an imperative programming style, similar to standard Python, making it intuitive for developers. While being user-friendly, Gluon can also achieve high performance by allowing models to be hybridized, which converts them into symbolic graphs for optimization.

Blocks: The Building Units of Gluon Models

gluon.Block is the fundamental building block for creating neural network layers and models in Gluon. Any neural network layer or a group of layers forming a component of a larger network is typically a subclass of Block.

Key characteristics of a Block:

• It can contain other Blocks (allowing for hierarchical model structures).
• It holds parameters (weights and biases) associated with the layers it defines.
• It defines a forward method (or hybrid_forward for HybridBlock) that specifies the computation performed by the block.

Example: A Simple Multi-Layer Perceptron (MLP) using nn.Sequential

gluon.nn.Sequential is a specialized Block that stacks other blocks in a sequence, where the output of one block is fed as input to the next.

import mxnet as mx
from mxnet import nd, gluon, autograd
from mxnet.gluon import nn # nn is a submodule of gluon containing predefined layers

mx.random.seed(0) # for reproducibility

# Dummy input data (batch_size=2, num_features=20)
x = nd.random.uniform(shape=(2, 20), ctx=mx.cpu())

# Define a sequential MLP
N_net = nn.Sequential()
with N_net.name_scope(): # Organizes parameter names
    N_net.add(nn.Dense(256, activation='relu')) # Fully connected layer with 256 units and ReLU activation
    N_net.add(nn.Dense(10))                     # Output layer with 10 units (e.g., for 10-class classification)

# Initialize parameters (weights and biases)
# This needs to be done before the first forward pass if not using deferred initialization with actual data.
# Alternatively, with deferred initialization (default for many layers), shapes are inferred on first forward pass.
N_net.initialize(ctx=mx.cpu()) # Initialize on CPU context

# Perform a forward pass
output = N_net(x)
print(output)

Output (actual values depend on random initialization):

<html>
 <body>
  <p>
   [[ 0.07308105  0.00406009 -0.00294911  0.0370597   0.00400224  0.008209
   -0.02912429 -0.05099454 -0.0040519   0.01902002]
   [ 0.09712769  0.00383658  0.01095401  0.02602063  0.02011784  0.00934159
   -0.01897359 -0.01587484 -0.01876913  0.02037334]]
   <NDArray 2x10 @cpu(0)>
  </p>
 </body>
</html>

Steps involved in using a Block (like nn.Sequential or custom blocks):

1. Define the Block: Instantiate it or define its architecture (e.g., by adding layers to nn.Sequential).
2. Initialize Parameters: Call .initialize() to set up the weights and biases. MXNet supports various initializers (e.g., Xavier, Normal). Shape inference often happens automatically on the first forward pass if not specified during initialization.
3. Forward Pass: Pass input data (NDArray) through the block by calling it like a function (e.g., N_net(x)). This executes the defined computation.
4. Backward Pass (during training): Use autograd.record() to record computations and then call .backward() on the loss scalar to compute gradients for the block’s parameters.

Understanding nn.Sequential Internally

We can implement a simplified version of Sequential to understand its core logic:

class MySequential(nn.Block):
    def __init__(self, **kwargs):
        super(MySequential, self).__init__(**kwargs)
        # self._children is an OrderedDict in the actual nn.Block to store sub-blocks

    def add(self, block):
        # In a real Block, self._children stores blocks with unique names.
        # For simplicity, we use a list here. Gluon uses an OrderedDict.
        if not hasattr(self, '_added_blocks'):
            self._added_blocks = []
        self._added_blocks.append(block)
        self.register_child(block) # Important for parameter collection etc.

    def forward(self, x):
        for block in self._added_blocks: # Iterate through added blocks
            x = block(x)
        return x

# Using MySequential
x_custom = nd.random.uniform(shape=(2, 20), ctx=mx.cpu())
N_net_custom = MySequential()
with N_net_custom.name_scope():
    N_net_custom.add(nn.Dense(256, activation='relu'))
    N_net_custom.add(nn.Dense(10))

N_net_custom.initialize(ctx=mx.cpu())
output_custom = N_net_custom(x_custom)
print(f"Output from MySequential (should be similar shape and type):\n{output_custom}")

Output (actual values will differ due to separate initialization):

<html>
 <body>
  <p>
Output from MySequential (should be similar shape and type):
[[ 0.01903114  0.02831168  0.01545286 -0.01109749  0.0350692  -0.02966933
  -0.01100099 -0.02197874 -0.01193464  0.00189657]
 [ 0.0183696   0.02339937  0.01078339 -0.0076901   0.03513341 -0.02404617
  -0.00755438 -0.01889096 -0.00740707  0.00199449]]
<NDArray 2x10 @cpu(0)>
  </p>
 </body>
</html>

Custom Blocks

For more complex architectures where layers are not simply sequential (e.g., networks with skip connections like ResNet, or multiple input/output branches), you need to create custom blocks by subclassing gluon.Block (or gluon.HybridBlock).

In a custom block, you typically:

1. Define layers as attributes in the __init__ method.
2. Implement the forward(self, x, ...) method to define the computation logic, using the layers defined in __init__.

class MLP(nn.Block):
    def __init__(self, **kwargs):
        super(MLP, self).__init__(**kwargs)
        # Define layers within the name_scope to manage parameter names
        with self.name_scope():
            self.hidden = nn.Dense(256, activation='relu')  # Hidden layer
            self.output_layer = nn.Dense(10)                # Output layer (renamed to avoid conflict with 'output' variable)

    def forward(self, x):
        # Define the forward computation
        hidden_activation = self.hidden(x)
        final_output = self.output_layer(hidden_activation)
        return final_output

x_mlp = nd.random.uniform(shape=(2, 20), ctx=mx.cpu())
N_net_mlp = MLP()
N_net_mlp.initialize(ctx=mx.cpu())
output_mlp = N_net_mlp(x_mlp)
print(f"Output from custom MLP block:\n{output_mlp}")

Output (values will depend on initialization):

<html>
 <body>
  <p>
Output from custom MLP block:
[[-0.02450891 -0.02139075 -0.04023137  0.01311899  0.00839001  0.01907403
   0.00343563  0.00571267 -0.02951591 -0.03287996]
 [-0.01569841 -0.01874819 -0.0320272   0.0182126   0.00212092  0.01575342
   0.0099024   0.00390883 -0.02429737 -0.02714972]]
<NDArray 2x10 @cpu(0)>
  </p>
 </body>
</html>

Custom Layers

While Gluon provides many common layers in gluon.nn, you might need to implement a custom layer from scratch. This is done by subclassing gluon.Block (for purely imperative layers) or gluon.HybridBlock (for layers that can be hybridized).

A Simple Custom Layer (without learnable parameters)

A custom layer needs to implement the forward method. If it has no learnable parameters, __init__ might just call the superclass constructor.

# A layer that normalizes input to the range [0, 1]
class NormalizationLayer(gluon.Block):
    def __init__(self, **kwargs):
        super(NormalizationLayer, self).__init__(**kwargs)

    def forward(self, x):
        # Ensure min_val is not equal to max_val to avoid division by zero
        min_val = nd.min(x)
        max_val = nd.max(x)
        if min_val == max_val:
            return nd.zeros_like(x)
        return (x - min_val) / (max_val - min_val)

x_norm_layer_input = nd.array([[1, 2, 3], [4, 5, 6]], ctx=mx.cpu(), dtype='float32')
norm_layer = NormalizationLayer()
# No .initialize() needed if there are no parameters, but good practice if it might gain params later.
# norm_layer.initialize(ctx=mx.cpu()) # Would be harmless
output_norm_layer = norm_layer(x_norm_layer_input)
print(f"Output from NormalizationLayer:\n{output_norm_layer}")

Output:

<html>
 <body>
  <p>
Output from NormalizationLayer:
[[0.  0.2 0.4]
 [0.6 0.8 1. ]]
<NDArray 2x3 @cpu(0)>
  </p>
 </body>
</html>

Hybridization: Combining Imperative and Symbolic Execution

gluon.HybridBlock is a special type of Block that can run in two modes: imperative (like a regular Block) and symbolic. After defining a HybridBlock, you can call its .hybridize() method. This attempts to convert the hybrid_forward method into a symbolic graph, which MXNet can then optimize for faster execution.

Benefits of hybridization:

• Flexibility: Develop and debug imperatively.
• Performance: Deploy with symbolic graph optimizations for speed and lower memory usage.
• Portability: Export the symbolic graph for inference in other language bindings or on devices.

When creating a HybridBlock, you implement hybrid_forward(self, F, x, ...) instead of forward. The F argument is crucial: it refers to either mxnet.ndarray (when running imperatively) or mxnet.symbol (when building the symbolic graph after .hybridize() is called). You must use F for all backend operations (e.g., F.relu(x) instead of nd.relu(x)).

Example: A Hybrid Normalization Layer

class NormalizationHybridLayer(gluon.HybridBlock):
    def __init__(self, **kwargs):
        super(NormalizationHybridLayer, self).__init__(**kwargs)

    def hybrid_forward(self, F, x):
        min_val = F.min(x)
        max_val = F.max(x)
        # Add a small epsilon to prevent division by zero if min_val == max_val
        # F.broadcast_xxx operations are common for N-D array compatibility
        denominator = F.broadcast_sub(max_val, min_val) + 1e-8
        return F.broadcast_div(F.broadcast_sub(x, min_val), denominator)

input_hybrid_norm = nd.array([1, 2, 3, 4, 5, 6], ctx=mx.cpu(), dtype='float32')
layer_hybd = NormalizationHybridLayer()
# layer_hybd.initialize() # Not strictly needed without params, but good practice

# Run imperatively first
output_imperative = layer_hybd(input_hybrid_norm)
print(f"Imperative output:\n{output_imperative}")

# Hybridize the layer
layer_hybd.hybridize()

# Run again (now using symbolic graph internally)
output_hybridized = layer_hybd(input_hybrid_norm)
print(f"Hybridized output:\n{output_hybridized}")

Output:

<html>
 <body>
  <p>
Imperative output:
[0.         0.2        0.4        0.6        0.8        1.        ]
<NDArray 6 @cpu(0)>
Hybridized output:
[0.         0.2        0.4        0.6        0.8        1.        ]
<NDArray 6 @cpu(0)>
  </p>
 </body>
</html>

Note: Hybridization is a compilation step to a symbolic graph. The choice of CPU or GPU for execution is separate and determined by the context of the input data and layer parameters.

Difference between Block and HybridBlock

The main difference is the ability to hybridize. Block always runs imperatively. HybridBlock can run imperatively before .hybridize() is called, and symbolically (compiled) afterwards. This requires using the F argument in hybrid_forward for all MXNet operations to ensure compatibility with both mx.nd and mx.sym backends.

Using Custom Layers in a Network

Custom layers, especially HybridBlocks, can be combined with standard layers using containers like gluon.nn.Sequential or gluon.nn.HybridSequential. HybridSequential is itself a HybridBlock and is used for stacking other HybridBlocks.

net = gluon.nn.HybridSequential()
with net.name_scope():
    net.add(nn.Dense(5, activation='relu')) # Standard dense layer
    net.add(NormalizationHybridLayer())   # Our custom hybrid layer
    net.add(nn.Dense(1))                  # Another standard dense layer

net.initialize(mx.init.Xavier(magnitude=2.24), ctx=mx.cpu())

# Hybridize the entire network
net.hybridize()

# Dummy input
input_data_for_net = nd.random.uniform(low=-10, high=10, shape=(10, 2), ctx=mx.cpu())
output_from_net = net(input_data_for_net)
print(f"Output from hybrid sequential net (shape {output_from_net.shape}):\n{output_from_net[0:3]}") # Print first 3 rows

Output (first 3 rows, actual values will vary):

<html>
 <body>
  <p>
Output from hybrid sequential net (shape (10, 1)):
[[-0.40172195]
 [-0.40172195]
 [-0.40172195]]
<NDArray 3x1 @cpu(0)>
  </p>
 </body>
</html>

Custom Layer Parameters

Layers often have learnable parameters (weights, biases) or fixed configuration parameters. In Gluon, parameters are managed using self.params, which is a ParameterDict object available in every Block.

• Learnable Parameters: Use self.params.get(...) to define them. Their shapes can often be inferred (allow_deferred_init=True). They are updated during training.
• Fixed Parameters/Buffers: Can also be stored if needed, sometimes as constants or non-differentiable parameters.

Example: Custom Layer with Learnable Weights and Fixed Scales

class ScaledFullyConnectedLayer(gluon.HybridBlock):
    def __init__(self, hidden_units, fixed_scales_data, **kwargs):
        super(ScaledFullyConnectedLayer, self).__init__(**kwargs)
        self.hidden_units = hidden_units
        with self.name_scope():
            # Learnable weights for a fully connected operation
            # Shape is (output_units, input_units_or_0_for_inference)
            self.weights = self.params.get('weights',
                                          shape=(hidden_units, 0), # Input dim inferred later
                                          allow_deferred_init=True,
                                          init=mx.init.Normal(0.01))
            # Fixed scales (not trainable)
            self.scales = self.params.get('scales',
                                         shape=fixed_scales_data.shape,
                                         init=mx.init.Constant(fixed_scales_data.asnumpy()),
                                         differentiable=False) # Not updated by optimizer
            # Bias (learnable)
            self.bias = self.params.get('bias',
                                       shape=(hidden_units,),
                                       init=mx.init.Constant(0),
                                       allow_deferred_init=True)

    def hybrid_forward(self, F, x, weights, scales, bias):
        # Fully connected operation using F for backend compatibility
        fc_out = F.FullyConnected(data=x, weight=weights, bias=bias, num_hidden=self.hidden_units, no_bias=False)
        # Element-wise multiplication with scales
        scaled_out = F.broadcast_mul(fc_out, scales)
        return scaled_out

# Example usage
num_features_in = 10
num_hidden_units = 5
fixed_scales = nd.array([0.1, 0.5, 1.0, 0.5, 0.1], ctx=mx.cpu()) # Scales for each hidden unit

scaled_fc_layer = ScaledFullyConnectedLayer(hidden_units=num_hidden_units, fixed_scales_data=fixed_scales)
scaled_fc_layer.initialize(ctx=mx.cpu())

dummy_input_fc = nd.random.normal(shape=(3, num_features_in), ctx=mx.cpu())
output_scaled_fc = scaled_fc_layer(dummy_input_fc)

print(f"Input shape: {dummy_input_fc.shape}")
print(f"Output shape from ScaledFullyConnectedLayer: {output_scaled_fc.shape}")
print(f"Parameters in layer: {scaled_fc_layer.collect_params()}")