Function

class dragon.vm.torch.autograd.Function(
  key,
  dev,
  **kwargs
)[source]

Dispatch the tensor operation.

Each tensor operation takes a function instance to dispatch the specific execution.

To define a new function class, overrides attributes and forward:

class MyReLU(autograd.Function):

    def __init__(self, key, dev, **kwargs):
        super(MyReLU, self).__init__(key, dev, **kwargs)
        self.alpha = kwargs.get('alpha', 0)

    def attributes(self):
        return {'op_type': 'Relu', 'arguments': {'alpha': float(self.alpha)}}

    def forward(self, input):
        outputs = [self.alloc()]
        return self.dispatch([input], outputs)

Function is executed by instantiating attributes and applying to inputs:

def my_relu(input, alpha=0):
    return MyReLU.instantiate(input.device, alpha=alpha).apply(input)

__init__

Function.__init__(
  key,
  dev,
  **kwargs
)[source]

Create a Function.

Parameters:

Methods

alloc

Function.alloc(out=None)[source]

Return or bind the executing device to output tensor.

Parameters:
Returns:

Union[dragon.vm.torch.device, dragon.vm.torch.Tensor] – The executing device or output tensor.

apply

Function.apply(
  *args,
  **kwargs
)[source]

Apply this function to inputs.

attributes

Function.attributes()[source]

Return the function attributes.

Returns:
dict – The attribute dict.

dispatch

Function.dispatch(
  inputs,
  outputs,
  no_grad=False,
  check_device=True,
  callback=None
)[source]

Dispatch the execution.

instantiate

classmethod Function.instantiate(
  device,
  **kwargs
)[source]

Return an instance of this function.

Parameters:
Returns:

dragon.vm.torch.autograd.Function – The function instance.

forward

Function.forward(
  *inputs,
  **kwargs
)[source]

Define the execution.