mindspore.nn.WithLossCell

class mindspore.nn.WithLossCell(backbone, loss_fn)[source]

Cell with loss function.

Wraps the network with loss function. This Cell accepts data and label as inputs and the computed loss will be returned.

Parameters

backbone (Cell) – The target network to wrap.
loss_fn (Cell) – The loss function used to compute loss.

Inputs:

data (Tensor) - Tensor of shape \((N, \ldots)\).
label (Tensor) - Tensor of shape \((N, \ldots)\).

Outputs:

Tensor, a scalar tensor with shape \(()\).

Raises: TypeError – If dtype of data or label is neither float16 nor float32.

Supported Platforms:: Ascend GPU

Examples

>>> net = Net()
>>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=False)
>>> net_with_criterion = nn.WithLossCell(net, loss_fn)
>>>
>>> batch_size = 2
>>> data = Tensor(np.ones([batch_size, 1, 32, 32]).astype(np.float32) * 0.01)
>>> label = Tensor(np.ones([batch_size, 10]).astype(np.float32))
>>>
>>> output_data = net_with_criterion(data, label)

property backbone_network

Returns the backbone network.

Returns: Cell, the backbone network.