Watch a model actually learn

The batch computes actual image loss and accumulated gradients.

export function evaluateTraining(records, data, withGradient=false) {
  let loss=0; const gradients=[0,0,0],frames=[];
  const count=data.reduce((n,frame)=>n+frame.samples.length,0);
  for(const frame of data) {
    const pixels=[];
    for(const sample of frame.samples) {
      const options={target:sample.target,finite:false};
      const position=withGradient?gradient(records,0,frame.time,sample.pixel,options):null;
      const result=position?.result??forward(records,frame.time,sample.pixel,options);
      loss+=result.loss/count;
      if(withGradient) {
        gradients[0]+=position.analytic/count;
        gradients[1]+=position.analytic*(frame.time-records[0][3])/count;
        gradients[2]+=gradient(records,20,frame.time,sample.pixel,options).analytic/count;
      }
      pixels.push(result);
    }
    frames.push(pixels);
  }
  return {loss,gradients,frames};
}

The update is simultaneous and uses pre-update gradients.

export function trainingIteration(records, data, rate=1) {
  if(!Number.isFinite(rate)||rate<=0)throw new RangeError('Learning rate must be positive');
  const beforeRecords=copyRecords(records),before=evaluateTraining(beforeRecords,data,true);
  const afterRecords=copyRecords(beforeRecords);
  // All gradients come from the same pre-update batch; commit simultaneously.
  trainingParameters.forEach((field,j)=>{afterRecords[0][field]-=rate*before.gradients[j];});
  const after=evaluateTraining(afterRecords,data);
  if(!Number.isFinite(after.loss)||!afterRecords.flat().every(Number.isFinite))throw new Error('Non-finite training state');
  return {beforeRecords,before,afterRecords,after,rate};

The separate Python experiment uses PyTorch autograd and Adam.

    optimizer = torch.optim.Adam(fitted, lr=0.025)
    history = []
    started = time.perf_counter()
    for step in range(args.steps):
        optimizer.zero_grad(set_to_none=True)
        loss = (render(*fitted, times, cameras)-truth).square().mean()
        if not torch.isfinite(loss):
            raise RuntimeError("Non-finite training loss")
        loss.backward()
        optimizer.step()