compute_norm_for_blobs.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from caffe2.python import core, schema, muji
from caffe2.python.modeling.net_modifier import NetModifier


import numpy as np


class ComputeNormForBlobs(NetModifier):
    """
    This class modifies the net passed in by adding ops to compute norms for
    certain blobs.

    Args:
        blobs: list of blobs to compute norm for
        logging_frequency: frequency for printing norms to logs
        p: type of norm. Currently it supports p=1 or p=2
        compute_averaged_norm: norm or averaged_norm (averaged_norm = norm/size)
    """

    def __init__(self, blobs, logging_frequency, p=2, compute_averaged_norm=False):
        self._blobs = blobs
        self._logging_frequency = logging_frequency
        self._p = p
        self._compute_averaged_norm = compute_averaged_norm
        self._field_name_suffix = '_l{}_norm'.format(p)
        if compute_averaged_norm:
            self._field_name_suffix = '_averaged' + self._field_name_suffix

    def modify_net(self, net, init_net=None, grad_map=None, blob_to_device=None,
                   modify_output_record=False):

        p = self._p
        compute_averaged_norm = self._compute_averaged_norm

        CPU = muji.OnCPU()
        # if given, blob_to_device is a map from blob to device_option
        blob_to_device = blob_to_device or {}
        for blob_name in self._blobs:
            blob = core.BlobReference(blob_name)
            if not net.BlobIsDefined(blob):
                raise Exception('blob {0} is not defined in net {1}'.format(
                    blob, net.Name()))
            if blob in blob_to_device:
                device = blob_to_device[blob]
            else:
                device = CPU

            with core.DeviceScope(device):
                norm_name = net.NextScopedBlob(prefix=blob + self._field_name_suffix)
                cast_blob = net.Cast(
                    blob,
                    net.NextScopedBlob(prefix=blob + '_float'),
                    to=core.DataType.FLOAT
                )
                norm = net.LpNorm(
                    cast_blob, norm_name, p=p, average=compute_averaged_norm
                )

                if self._logging_frequency >= 1:
                    net.Print(norm, [], every_n=self._logging_frequency)

                if modify_output_record:
                    output_field_name = str(blob) + self._field_name_suffix
                    output_scalar = schema.Scalar((np.float, (1,)), norm)

                    if net.output_record() is None:
                        net.set_output_record(
                            schema.Struct((output_field_name, output_scalar))
                        )
                    else:
                        net.AppendOutputRecordField(
                            output_field_name,
                            output_scalar)

    def field_name_suffix(self):
        return self._field_name_suffix