metrics.py

# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Loss and accuracy."""
from mindspore import nn
from mindspore import Tensor
from mindspore.common import dtype as mstype
from mindspore.ops import operations as P
from mindspore.common.parameter import ParameterTuple
from mindspore.ops import composite as C
from mindspore.ops import functional as F


class Loss(nn.Cell):
    """Softmax cross-entropy loss with masking."""
    def __init__(self, label, mask, weight_decay, param):
        super(Loss, self).__init__()
        self.label = Tensor(label)
        self.mask = Tensor(mask)
        self.loss = P.SoftmaxCrossEntropyWithLogits()
        self.one = Tensor(1.0, mstype.float32)
        self.zero = Tensor(0.0, mstype.float32)
        self.mean = P.ReduceMean()
        self.cast = P.Cast()
        self.l2_loss = P.L2Loss()
        self.reduce_sum = P.ReduceSum()
        self.weight_decay = weight_decay
        self.param = param

    def construct(self, preds):
        """Calculate loss"""
        param = self.l2_loss(self.param)
        loss = self.weight_decay * param
        preds = self.cast(preds, mstype.float32)
        loss = loss + self.loss(preds, self.label)[0]
        mask = self.cast(self.mask, mstype.float32)
        mask_reduce = self.mean(mask)
        mask = mask / mask_reduce
        loss = loss * mask
        loss = self.mean(loss)
        return loss


class Accuracy(nn.Cell):
    """Accuracy with masking."""
    def __init__(self, label, mask):
        super(Accuracy, self).__init__()
        self.label = Tensor(label)
        self.mask = Tensor(mask)
        self.equal = P.Equal()
        self.argmax = P.Argmax()
        self.cast = P.Cast()
        self.mean = P.ReduceMean()

    def construct(self, preds):
        preds = self.cast(preds, mstype.float32)
        correct_prediction = self.equal(self.argmax(preds), self.argmax(self.label))
        accuracy_all = self.cast(correct_prediction, mstype.float32)
        mask = self.cast(self.mask, mstype.float32)
        mask_reduce = self.mean(mask)
        mask = mask / mask_reduce
        accuracy_all *= mask
        return self.mean(accuracy_all)


class LossAccuracyWrapper(nn.Cell):
    """
    Wraps the GCN model with loss and accuracy cell.

    Args:
        network (Cell): GCN network.
        label (numpy.ndarray): Dataset labels.
        mask (numpy.ndarray): Mask for training, evaluation or test.
        weight_decay (float): Weight decay parameter for weight of the first convolution layer.
    """

    def __init__(self, network, label, mask, weight_decay):
        super(LossAccuracyWrapper, self).__init__()
        self.network = network
        self.loss = Loss(label, mask, weight_decay, network.trainable_params()[0])
        self.accuracy = Accuracy(label, mask)

    def construct(self):
        preds = self.network()
        loss = self.loss(preds)
        accuracy = self.accuracy(preds)
        return loss, accuracy


class LossWrapper(nn.Cell):
    """
    Wraps the GCN model with loss.

    Args:
        network (Cell): GCN network.
        label (numpy.ndarray): Dataset labels.
        mask (numpy.ndarray): Mask for training.
        weight_decay (float): Weight decay parameter for weight of the first convolution layer.
    """

    def __init__(self, network, label, mask, weight_decay):
        super(LossWrapper, self).__init__()
        self.network = network
        self.loss = Loss(label, mask, weight_decay, network.trainable_params()[0])

    def construct(self):
        preds = self.network()
        loss = self.loss(preds)
        return loss


class TrainOneStepCell(nn.Cell):
    r"""
    Network training package class.

    Wraps the network with an optimizer. The resulting Cell be trained without inputs.
    Backward graph will be created in the construct function to do parameter updating. Different
    parallel modes are available to run the training.

    Args:
        network (Cell): The training network.
        optimizer (Cell): Optimizer for updating the weights.
        sens (Number): The scaling number to be filled as the input of backpropagation. Default value is 1.0.

    Outputs:
        Tensor, a scalar Tensor with shape :math:`()`.

    Examples:
        >>> net = Net()
        >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
        >>> optim = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
        >>> loss_net = nn.WithLossCell(net, loss_fn)
        >>> train_net = nn.TrainOneStepCell(loss_net, optim)
    """

    def __init__(self, network, optimizer, sens=1.0):
        super(TrainOneStepCell, self).__init__(auto_prefix=False)
        self.network = network
        self.network.set_grad()
        self.network.add_flags(defer_inline=True)
        self.weights = ParameterTuple(network.trainable_params())
        self.optimizer = optimizer
        self.grad = C.GradOperation(get_by_list=True, sens_param=True)
        self.sens = sens

    def construct(self):
        weights = self.weights
        loss = self.network()
        sens = P.Fill()(P.DType()(loss), P.Shape()(loss), self.sens)
        grads = self.grad(self.network, weights)(sens)
        return F.depend(loss, self.optimizer(grads))


class TrainNetWrapper(nn.Cell):
    """
    Wraps the GCN model with optimizer.

    Args:
        network (Cell): GCN network.
        label (numpy.ndarray): Dataset labels.
        mask (numpy.ndarray): Mask for training, evaluation or test.
        config (ConfigGCN): Configuration for GCN.
    """

    def __init__(self, network, label, mask, config):
        super(TrainNetWrapper, self).__init__(auto_prefix=True)
        self.network = network
        loss_net = LossWrapper(network, label, mask, config.weight_decay)
        optimizer = nn.Adam(loss_net.trainable_params(),
                            learning_rate=config.learning_rate)
        self.loss_train_net = TrainOneStepCell(loss_net, optimizer)
        self.accuracy = Accuracy(label, mask)

    def construct(self):
        loss = self.loss_train_net()
        accuracy = self.accuracy(self.network())
        return loss, accuracy