Apache MXNet - Python API autograd and initializer

MXNet Python API: Autograd and Initializer

This chapter delves into the specifics of MXNet’s autograd API for automatic differentiation and the initializer API for setting initial parameter values in neural networks.

mxnet.autograd

The mxnet.autograd module provides tools for automatic differentiation on NDArrays.

Class: mxnet.autograd.Function

mxnet.autograd.Function allows users to define custom differentiable operations. If you need non-standard gradient computations not covered by MXNet’s built-in operators, you can subclass Function and implement custom forward and backward methods.

Workflow for custom autograd Functions:

• Define the forward computation in the forward method. This method takes input NDArrays and returns output NDArrays.
• Implement the custom gradient calculation in the backward method. This method receives the gradient of the subsequent function’s output (often denoted dy or head_grads) and must return gradients with respect to each input of the forward method.
• During backpropagation, mxnet.autograd will invoke your custom backward method instead of relying on the chain rule for predefined ops.
• You can use self.save_for_backward(*tensors) in forward to cache any tensors needed for the backward computation, retrievable via self.saved_tensors in backward.

Example: Defining a stable sigmoid function with custom forward and backward passes.

import mxnet as mx
from mxnet import nd, autograd

class StableSigmoid(autograd.Function):
    def forward(self, x):
        # Numerically stable sigmoid: y = 1 / (1 + exp(-x))
        y = 1 / (1 + nd.exp(-x))
        self.save_for_backward(y)  # Save y for use in backward pass
        return y

    def backward(self, dy):
        # Gradient of sigmoid: dy * y * (1 - y)
        y, = self.saved_tensors  # Retrieve saved y
        return dy * y * (1 - y)

# Using the custom function
custom_sigmoid_op = StableSigmoid()

x_input = nd.random.uniform(shape=(5,))
x_input.attach_grad() # Mark x_input to compute its gradient

with autograd.record():
    output_y = custom_sigmoid_op(x_input)

# Backpropagate from a dummy head gradient (e.g., ones)
output_y.backward(nd.ones_like(output_y))

dx_gradient = x_input.grad

print("Input x:")
print(x_input.asnumpy())
print("\nOutput y (sigmoid(x)):")
print(output_y.asnumpy())
print("\nGradient dx (custom):")
print(dx_gradient.asnumpy())

Output (values will vary due to random input):

Input x:
[0.6288061  0.03209561 0.9043807  0.07945406 0.3247892 ]

Output y (sigmoid(x)):
[0.6521595  0.5080225  0.711851   0.5198517  0.58049846]

Gradient dx (custom):
[0.22686191 0.24993561 0.20499429 0.2496056  0.24351722]

Key mxnet.autograd Methods

Commonly used functions and context managers within mxnet.autograd:

Method/Context Manager	Description
record(train_mode=True)	Returns a context manager. Operations within this scope are recorded for gradient computation. train_mode affects behavior of layers like Dropout, BatchNorm.
pause(train_mode=True)	Returns a context manager. Operations within this scope are NOT recorded. Useful for parts of computation that shouldn’t affect gradients (e.g., evaluation metrics within a training loop).
is_recording()	Returns True if currently inside a record() scope, False otherwise.
is_training()	Returns True if currently in training mode (affects layers like Dropout, BatchNorm), False otherwise.
mark_variables(variables, gradients, grad_reqs='write')	Marks NDArrays as variables for gradient computation. Similar to ndarray.attach_grad() but offers more control (e.g., setting initial gradient values).
grad(heads, variables, head_grads=None, retain_graph=False, create_graph=False)	Computes gradients of heads (output NDArrays) with respect to variables (input NDArrays). Returns gradients as new NDArrays instead of storing them in .grad attributes.

Implementation Example: autograd.grad()

Using mxnet.autograd.grad() to compute gradients directly:

x = nd.ones((2,))
x.attach_grad() # Need to attach grad to x if it's a leaf variable we need gradient for

with autograd.record():
    # y = x * x
    # z = y + x
    # Let's do something like z = exp(x) + x
    z = nd.exp(x) + x

# Compute dz/dx
dx_grads_list = autograd.grad(heads=z, variables=[x], create_graph=False)
dx_grad = dx_grads_list[0]

print("x:", x.asnumpy())
print("z:", z.asnumpy())
print("dz/dx:", dx_grad.asnumpy()) # Should be exp(x) + 1. For x=1, e^1+1 approx 2.718+1 = 3.718

Output:

x: [1. 1.]
z: [3.7182817 3.7182817]
dz/dx: [3.7182817 3.7182817]

Using predict_mode() (or train_mode=False in record or pause) for inference sections:

# model = SomeGluonModel()
# x_input = nd.random.normal(shape=(1,10))
# with autograd.record(train_mode=True): # Training mode for parts like dropout/batchnorm
#     intermediate_y = model.feature_extractor(x_input)
# with autograd.predict_mode(): # Inference mode for parts not needing grad or special train behavior
#     final_output_logits = model.classifier(intermediate_y)
# # Or if only prediction is needed:
# # with autograd.predict_mode():
# #    predictions = model(x_input)
# # loss.backward() # if loss was computed in record() scope
print("predict_mode() is used to ensure parts of network run in inference mode.")

mxnet.initializer (or mx.init)

The mxnet.initializer module (often accessed via mx.init) provides various schemes for initializing the weights and biases of neural network layers. Proper initialization is crucial for effective training.

Common Initializer Classes

Key initializer classes and their purpose:

Initializer Class	Description
mx.init.Normal(sigma=0.01)	Initializes weights with random values from a normal distribution with mean 0 and specified standard deviation sigma.
mx.init.Uniform(scale=0.07)	Initializes weights with random values from a uniform distribution within [-scale, scale].
mx.init.Xavier(rnd_type='uniform', factor_type='avg', magnitude=3)	Implements Xavier initialization, designed to keep activation variances consistent across layers. rnd_type can be ‘gaussian’ or ‘uniform’. factor_type (‘avg’, ‘in’, ‘out’) determines how fan-in/fan-out are used.
mx.init.MSRAPrelu (or Kaiming)	Implements Kaiming He initialization, particularly effective for ReLU-like activations. Similar options to Xavier regarding distribution and fan type.
mx.init.Constant(value)	Initializes weights to a fixed scalar value or a specific NDArray.
mx.init.Zero()	Initializes weights to zero. Typically used for biases if not learned, or in specific scenarios.
mx.init.One()	Initializes weights to one.
mx.init.Bilinear()	Initializes weights for upsampling layers (e.g., in deconvolution) to perform bilinear interpolation.
mx.init.Mixed(patterns, initializers)	Applies different initializers to parameters based on their names matching specified regex patterns.

Implementation Example: Using Initializers

Applying an initializer when defining a Gluon layer or model:

from mxnet import gluon, init

# Example 1: Initialize a Dense layer with Xavier
layer_xavier = gluon.nn.Dense(units=64, weight_initializer=init.Xavier(), bias_initializer=init.Zero())
layer_xavier.initialize() # Actual initialization happens here
# print("Xavier initialized weight sample:", layer_xavier.weight.data()[0,0:5].asnumpy())

# Example 2: Using Mixed initializer for a Sequential model
net = gluon.nn.Sequential()
with net.name_scope():
    net.add(gluon.nn.Dense(256, activation='relu')) # Default init will apply if not specified in net.initialize
    net.add(gluon.nn.Dense(10))

# Define mixed initializers: biases to zero, weights with Normal(0.02)
mixed_init = init.Mixed(
    ['.*bias', '.*gamma', '.*beta', '.*running_mean', '.*running_var'], # Patterns for biases, batchnorm scales/shifts, running stats
    [init.Zero(), init.One(), init.Zero(), init.Zero(), init.One()] # Corresponding initializers
)
# Fallback for other weights (e.g., '.*weight') can be specified or use default.
# To specify for weights: net.initialize(init.Normal(0.02), ctx=mx.cpu())
# For mixed: net.initialize(mixed_init, ctx=mx.cpu())
# A more common mixed init pattern:
net.initialize(init.Mixed(['.*weight', '.*bias'],
                           [init.Normal(0.02), init.Zero()]),
               ctx=mx.cpu())

print("Initialized net[0].bias (Dense layer 0 bias):")
print(net[0].bias.data().asnumpy())
print("\nSample of net[0].weight (Dense layer 0 weight) - first 5 values of first row:")
print(net[0].weight.data()[0,:5].asnumpy())

Output (bias should be zeros, weights will be small random numbers):

Initialized net[0].bias (Dense layer 0 bias):
[0. 0. 0. ... 0. 0. 0.] (Shape will be (256,))

Sample of net[0].weight (Dense layer 0 weight) - first 5 values of first row:
[-0.00890991  0.01314114 -0.00038932  0.02028422  0.0158705 ]