Module hooks

Concise Implementation of Linear Regression

Framework linear regression

The same model, same data, same training — using the framework’s high-level layers and built-in losses:

• Model: one LazyLinear (or equivalent) instead of hand-rolled w, b.
• Loss: built-in MSELoss (no factor of ½).
• Optimizer: built-in SGD.

End result: ~5 lines of model code instead of 30. Same convergence on synthetic data.

Model setup

Wrap a single linear layer with the right output dimension. The “lazy” variant defers the input-dim shape until the first forward:

from d2l import mxnet as d2l
from mxnet import autograd, gluon, init, np, npx
from mxnet.gluon import nn
npx.set_np()

class LinearRegression(d2l.Module):
    """The linear regression model implemented with high-level APIs."""
    def __init__(self, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Dense(1)
        self.net.initialize(init.Normal(sigma=0.01))

Hook the layer into our Module interface (forward, configure_optimizers):

@d2l.add_to_class(LinearRegression)
def forward(self, X):
    return self.net(X)

Loss and optimizer

Built-in MSE — note it omits the 1/2 factor we used by hand:

@d2l.add_to_class(LinearRegression)
def loss(self, y_hat, y):
    fn = gluon.loss.L2Loss()
    return 2 * fn(y_hat, y).mean()  # Gluon's L2Loss includes 1/2; multiply by 2 to get plain MSE

Same SGD, instantiated with one call:

@d2l.add_to_class(LinearRegression)
def configure_optimizers(self):
    return gluon.Trainer(self.collect_params(),
                         'sgd', {'learning_rate': self.lr})

Train

Identical loop — the Trainer doesn’t care that the model is now a thin wrapper around a built-in layer:

model = LinearRegression(lr=0.03)
data = d2l.SyntheticRegressionData(w=d2l.tensor([2, -3.4]), b=4.2)
trainer = d2l.Trainer(max_epochs=3)
trainer.fit(model, data)

Compare with ground truth

Pull weights and bias back out of the layer:

@d2l.add_to_class(LinearRegression)
def get_w_b(self):
    return (self.net.weight.data(), self.net.bias.data())
w, b = model.get_w_b()

print(f'error in estimating w: {data.w - d2l.reshape(w, data.w.shape)}')
print(f'error in estimating b: {data.b - b}')

Errors are tiny — same recovery as the from-scratch version, less glue code.

Recap

• From scratch taught us what was happening; concise is what we’ll actually use.
• The high-level layers / losses / optimizers compose with the same Module / Trainer scaffold.
• Same minibatch loop, same convergence; one-line layer instead of hand-rolled parameters.