[Feature Request] nn.Module should also get a `device` attribute

[nehz]

Similar to the changes for tensors in 0.4, Modules should also get a .device attribute to allow easy creation of correct tensors.

Use case:

class Net(Module):
    def forward(self, x):
        x = x.to(self.device)

net = Net().cuda()
net.forward(Tensor([1]))   # Don't need conditional or additional .to()

[apaszke]

That’s not possible. Modules can hold parameters of different types on different devices, and so it’s not always possible to unambiguously determine the device.

[nehz]

Potentially a new mixin class that requires all parameters to be on same device and move .cuda() and .to() methods to that.
As in the case with Modules with parameters on difference devices, doing a .cuda() or .to() can also mess it all up.

E.G

class Module:
    # Remove .cuda(), .to(), or put warning
    ...

class SingleDevice:
    def to():
        # code from Module
        ...

    @property
    def device():
         ...

class Linear(Module, SingleDevice):
    ...
    pass

class Container(Module):
    ...
    pass

Container([net1, net2]).cuda() # <--- doesn't make sense if some modules on different devices
Linear(5, 5).cuda() # OK

Thoughts ?

[apaszke]

I think that's a bit too convoluted for such a simple thing. If you know all parameters are on a single device you can always do next(model.parameters()).device.

[acgtyrant]

How about make the device of nn.Module as not implemented? Then all the official implemented module inherited from nn.Module should have the uniform device for their parameters (if I am wrong, forget it) so that they can have the device attribute, so as to DataParallel and DistributedParallel while their device is their module's device.

So if the user define a new Module class, then he should implement the device attributes by himself, or get the unimplemented exception when he use the unimplemented device attribute.

[apaszke]

I can't see how a missing attribute error is less informative than a NonImplementedError. It's not like there's anything in the Module contract that forbids you from adding a .device property. You can always roll it on your own any time.

[spott]

In simple cases (where all parameters are on the same device, which is probably a pretty common case), .device could return the expected answer. In complicated situations, it could throw a value error or similar.

How expensive would it be to run device on all parameters() in a module, and then check if they are all the same?

[aliutkus]

allowing the parameters of a model to be on several devices is nice, but it's really painful not to have the .device accessible whenever they're all on the same. I guess it's the overwhelming majority of cases..

[cifkao]

Something like

@property
def device(self):
    devices = ({param.device for param in self.parameters()} |
               {buf.device for buf in self.buffers()})
    if len(devices) != 1:
        raise RuntimeError('Cannot determine device: {} different devices found'
                           .format(len(devices)))
    return next(iter(devices))

would be nice.

It could be a cached property, invalidated by any call to _apply() (which is used by to() et al.).

[minhlab]

This is an one-liner that will throw an exception if the parameters are on different devices.

device, = list(set(p.device for p in model.parameters()))