%matplotlib inline
from d2l import jax as d2l
import jax
from jax import numpy as jnp
from flax import nnx
import optax
import numpy as np
from PIL import Image
def cls_predictor(num_inputs, num_anchors, num_classes, rngs=None):
rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
return nnx.Conv(num_inputs, num_anchors * (num_classes + 1),
kernel_size=(3, 3), padding='same', rngs=rngs)Single Shot Detection
Single Shot MultiBox Detector (Liu et al., 2016) is the prototype single-stage detector. One forward pass produces class scores and box offsets for every anchor at every scale; NMS keeps the survivors.
The architecture: a CNN trunk, then a pyramid of feature maps at decreasing resolutions. Each level has its own pair of 1×1-style heads — one for class scores, one for box offsets. Predictions from all levels are concatenated.
SSD = base network + several multiscale feature blocks; each block has its own anchor predictor.
Training
Standard SGD loop, two evaluation metrics (class accuracy, box mean abs error). Read them together: class accuracy is dominated by many background anchors, while box error only makes sense on matched positive anchors.
@nnx.jit
def train_step(net, optimizer, X, Y):
def loss_fn(net):
anchors, cls_preds, bbox_preds = net(X)
bbox_labels, bbox_masks, cls_labels = d2l.multibox_target(anchors, Y)
l = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels,
bbox_masks)
return l.mean(), (cls_preds, bbox_preds, bbox_labels, bbox_masks,
cls_labels)
(loss, aux), grads = nnx.value_and_grad(loss_fn, has_aux=True)(net)
optimizer.update(net, grads)
cls_preds, bbox_preds, bbox_labels, bbox_masks, cls_labels = aux
# Compute scalar metrics inside the jit so we don't ship large tensors
# back to the host every step.
cls_correct = (cls_preds.argmax(axis=-1).astype(cls_labels.dtype)
== cls_labels).sum()
cls_count = jnp.array(cls_labels.size, dtype=cls_correct.dtype)
bbox_abs_sum = jnp.abs((bbox_labels - bbox_preds) * bbox_masks).sum()
bbox_count = jnp.array(bbox_labels.size, dtype=bbox_abs_sum.dtype)
return loss, cls_correct, cls_count, bbox_abs_sum, bbox_count
num_epochs, timer = 20, d2l.Timer()
animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],
legend=['class error', 'bbox mae'])
for epoch in range(num_epochs):
# Sum of training accuracy, no. of examples in sum of training accuracy,
# Sum of absolute error, no. of examples in sum of absolute error
cls_correct_sum = 0.0
cls_total = 0.0
bbox_abs_total = 0.0
bbox_total = 0.0
for features, target in train_iter:
timer.start()
X, Y = jnp.asarray(features), jnp.asarray(target)
loss, cls_correct, cls_count, bbox_abs_sum, bbox_count = train_step(
train_net, optimizer, X, Y)
cls_correct_sum += cls_correct
cls_total += cls_count
bbox_abs_total += bbox_abs_sum
bbox_total += bbox_count
cls_err = float(1 - cls_correct_sum / cls_total)
bbox_mae = float(bbox_abs_total / bbox_total)
animator.add(epoch + 1, (cls_err, bbox_mae))
print(f'class err {cls_err:.2e}, bbox mae {bbox_mae:.2e}')
print(f'{len(train_iter) * batch_size / timer.stop():.1f} examples/sec on '
f'{str(jax.devices()[0])}')
class err 3.53e-03, bbox mae 3.52e-03
8036.9 examples/sec on cuda:0Detect bananas
Visualize all predictions with confidence ≥ 0.9. The useful thing to notice is not the raw tensor length, but whether NMS leaves one tight high-confidence box per banana:
def display(img, output, threshold):
d2l.set_figsize((5, 5))
fig = d2l.plt.imshow(img)
for row in output:
score = float(row[1])
if score < threshold:
continue
h, w = img.shape[:2]
bbox = [row[2:6] * jnp.array((w, h, w, h))]
d2l.show_bboxes(fig.axes, bbox, '%.2f' % score, 'w')
display(img, output, threshold=0.9)
def smooth_l1(data, scalar):
out = []
for i in data:
if abs(i) < 1 / (scalar ** 2):
out.append(((scalar * i) ** 2) / 2)
else:
out.append(abs(i) - 0.5 / (scalar ** 2))
return jnp.array(out)
sigmas = [10, 1, 0.5]
lines = ['-', '--', '-.']
x = jnp.arange(-2, 2, 0.1)
d2l.set_figsize()
for l, s in zip(lines, sigmas):
y = smooth_l1(x, scalar=s)
d2l.plt.plot(x, y, l, label='sigma=%.1f' % s)
d2l.plt.legend();
def focal_loss(gamma, x):
return -(1 - x) ** gamma * jnp.log(x)
x = jnp.arange(0.01, 1, 0.01)
for l, gamma in zip(lines, [0, 1, 5]):
y = d2l.plt.plot(x, focal_loss(gamma, x), l, label='gamma=%.1f' % gamma)
d2l.plt.legend();