Sqrt Scaling Inferred Batch

Problem

Invert the square-root batch scaling rule to infer a batch size.

Given a base learning rate and base batch, return base_batch * (lr_new / base_lr)^2, the batch size that would justify lr_new under sqrt scaling. If base_lr is not positive, return 0.0.

Hint

Squaring the learning-rate ratio recovers the batch-size ratio.

Test cases

case 1

in: [0.2,0.1,32]

expected: 128

case 2

in: [0.05,0.1,32]

expected: 8

case 3

in: [0.1,0.1,64]

expected: 64

case 4

in: [0,0,8]

expected: 0

Starter codePython

def sqrt_scaling_batch(lr_new, base_lr, base_batch):
    # Your code here
    pass

4 test cases · run your code against them in the editor