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.