split_and_load can now handle num_ctx > num_data. Github Issue #13909

[mightydeveloper]

Description

Handles Issue #13909
when last batch size is smaller than the number of contexts,
the previous utility function gluon.utils.split_and_load throws an exception ValueError: Too many slices for data with shape ....
However, we can just put one data per context and ignore the remaining contexts.
This integrates nicely with the given reproducing example code. (User does not need to modify the code)

for d, l in zip(data, label):
    with autograd.record():
        out = net(d)
        losses.append(loss_fn(out, l))
for loss in losses:
    loss.backward()

where data and label will only output data that exists in the first few of the given context.
the forward / backward pass is only done in the contexts where needed. (remaining contexts does not need to do a fake forward/backward pass)
My concern is people can make mistakes when calculating mean of the losses.

So it would be nice if we add this behavior to the documentation. (split_and_load can output a list that has a size less than number of contexts when even_split=False)

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.
-> I believe this PR is just tiny change

[mightydeveloper]

Just an FYI, I made a misleading comment on the code previously, so I fixed it.

[wkcn]

I am not sure how the weights update for the context without inpus when backward.

[mightydeveloper]

@wkcn
In the example above, losses only contains loss terms that are forwarded by real data samples. (meaning that unnecessary(possibly fake) context loss are not appended from the previous for loop.)
So I believe when we do loss.backward(), only the contexts that are appended from the previous for loop will run and calculate gradients for previously marked variables from previous for loop.

For example, suppose we have 3 examples left from the dataset and have 5 GPU contexts.
We would only call losses.append(loss_fn(out, l)) 3 times, marking variables in 0, 1, 2 GPU contexts.
So when we call

for loss in losses:
    loss.backward()

only gradients for GPU 0, 1, 2 will be calculated and when we eventually call trainer.step(), the weights will be updated.
Does this answer help you? (I might have misunderstood your concern)

(+ Also, I noticed that I have an indentation error in my original PR request description so I just edited it.)

[wkcn]

@mightydeveloper
Thank you for the detailed explanation!
Yes, only gradients for GPU 0, 1, 2 will be calculated.

However, it seems that trainer.step() will update the weights by using all gradients on 5 GPU contexts.
The gradients on GPU 3, 4 may be not zero-tensors.

The code of trainer.step()

[wkcn]

LGTM : )
Thanks for your contribution!

[mightydeveloper]

However, it seems that trainer.step() will update the weights by using all gradients on 5 GPU contexts.
The gradients on GPU 3, 4 may be not zero-tensors.

Thanks for checking and providing relevant code link!

So... If I want to zero out gradients, I guess I should either

1. call trainer._params.zero_grad()
2. or just net.collect_params().zero_grad()
   (Assuming that I initialized trainer with Trainer(net.collect_params(), ...)
3. or maybe call trainer.step(ignore_stale_grad=True) instead?

Would all three of them work? which one would be better?

[wkcn]

calling trainer.step(ignore_stale_grad=True) is okay : )

[piyushghai]

@mightydeveloper Thanks for your contributions.

@mxnet-label-bot Add [Gluon]

[wkcn]

Merged. Thanks for your contribution!