Added a couple of functions to manage distributed groups with `InplaceABNSync`:
* `active_group`: create a distributed group where each worker can decide wether to participate or not.
* `set_active_group`: scan a model, passing a distributed group to all layers that implement a `set_group()` method.
These are intended to simplify handling of asymmetric computational graphs in `DistributedDataParallel` when using `InplaceABNSync`. A typical usage is as follows:
python
class DynamicModel(nn.Module):
def __init__(self):
super(DynamicModel, self).__init__()
self.conv1 = nn.Conv2d(4, 4, 1)
self.bn1 = InplaceABNSync(4)
self.conv2 = nn.Conv2d(4, 4, 1)
self.bn2 = InplaceABNSync(4)
def forward(x):
x = self.conv1(x)
x = self.bn1(x)
Call some data-dependent function telling us wether the second part of the network
should be traversed or not
active = self.get_active(x)
Create process group containing only the active workers, pass it to bn2
set_active_group(self.bn2, active_group(active))
Run the second part of the network only if active is True
if active:
x = self.conv2(x)
x = self.bn2(x)
return x