Model Checkpointing

A checkpoint records state that you can load later. In this exercise, serialize a small snapshot containing a weight vector, a training-step counter and a learning rate. Return the JSON text; do not write a file.

Your task

Implement serialize_state(weights, step, lr):

weights is a one-dimensional floating-point tensor of shape (d,).
step is a nonnegative integer.
lr is a finite, nonnegative learning rate.

Weights contain finite values. Convert them to a Python list of numbers, preserving their order and values. Build a dictionary with these keys in this order:

weights: the list of weights
step: a Python int
lr: a Python float

Use the standard-library json module to serialize that dictionary and return a Python string. Preserve the input tensor. Do not use torch.save, pickle, a tensor's printed representation or manual string concatenation.

Exact output format

This exercise compares the returned string exactly. Use the default json.dumps formatting, with spaces after commas and colons, no indentation, and no sorted keys. Do not round the weight values or append a newline.

For weights = [1.0, -0.5], step = 12 and lr = 0.01, return:

{"weights": [1.0, -0.5], "step": 12, "lr": 0.01}

The return value is the JSON string itself, not the dictionary and not a second JSON encoding of that string. JSON objects can generally use other key orders or spacing, but those alternatives are outside this exercise's exact-string contract.

What this snapshot preserves

json.loads can recover the list, step and learning rate. Reconstructing a tensor also requires choosing a framework, dtype and device; none of those is recorded in this format. Decimal output may show the precision of the stored tensor values, so serialize the values you have rather than rounding them to their originally typed spelling.

This is not a complete training-resumption checkpoint. Depending on the training setup, resuming the same computation can also require optimizer buffers, scheduler state, random-number generator state and the position in the data stream. Saving just weights, step and learning rate does not preserve those details.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Save enough state for the intended restart

A weights file can support inference when the model architecture and required metadata are available. Continuing training usually needs more: optimizer moments, step counters, learning-rate schedule state, and sometimes random and data-loader state.

Suppose a weight is 2 at step 10. Restoring that weight does not reproduce step 11 if Adam's moment estimates or the next batch differ. Decide whether the task requires loading weights or reproducing a training continuation.

Verify a serialization round trip

A small JSON checkpoint exercise converts arrays and scalars into serializable values, then reconstructs them. JSON does not automatically preserve tensor dtype, device, or all shape information. Those must be encoded if the format requires them; follow the exercise's smaller contract when they are intentionally omitted.

Check the restored values and types. Missing fields, wrong dimensions, or an unsupported format version should produce a visible failure, not a successful restore with unrelated defaults. JSON spacing and object-key order usually do not affect the restored values, unless the task explicitly requires a particular serialized form.

Keep the best model separate from the live model

When recording the best validation checkpoint in memory, make an independent copy. Keeping a reference to mutable live weights can cause the saved version to change during later updates.

This kind of checkpoint differs from activation checkpointing, which saves training memory by recomputing intermediate activations during backpropagation. Activation checkpointing does not by itself save a model for restarting later.

Detecting runtime…