sciai.architecture.FirstOutputCell

View Source On Gitee
class sciai.architecture.FirstOutputCell(backbone)[source]

Network that return the first output of given network.

Parameters

backbone (Callable) – Original network.

Inputs:
  • *inputs (Tensor) - Original network inputs.

Outputs:

Union(Tensor, tuple[Tensor]), The first output of the original network.

Supported Platforms:

GPU CPU Ascend

Examples

>>> import mindspore as ms
>>> from mindspore import ops, nn
>>> from sciai.architecture.basic_block import FirstOutputCell
>>> class Net2In3Out(nn.Cell):
>>>     def __init__(self):
>>>         super().__init__()
>>>     def construct(self, x, y):
>>>         out1 = x + y
>>>         out2 = 2 * x + y
>>>         out3 = x * x + 4 * y * y + 3 * y
>>>         return out1.sum(), out2.sum(), out3.sum()
>>> net = Net2In3Out()
>>> first_output_cell = FirstOutputCell(net)
>>> x, y = ops.ones((2, 3), ms.float32), ops.ones((2, 3), ms.float32)
>>> res = first_output_cell(x, y)
>>> print(res)
12.0