ContrastiveLossLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers
{
    public class ContrastiveLossLayer<T> : LossLayer<T>
    {
        Blob<T> m_blobDiff; // cached for backward pass.
        Blob<T> m_blobDistSq; // cached for backward pass.
        Blob<T> m_blobDiffSq; // cached for backward pass.
        Blob<T> m_blobSummerVec; // tmp storage for gpu forward pass.
        Blob<T> m_blobSimilar; // tmp storage for backward pass.
        Blob<T> m_blobDistScale; // tmp storage for distance scale applied ot matching distances (optional).
        Blob<T> m_blobPrimary; // target of similar, or dissimilar image.
        T[] m_rgMatches = null;
        int m_nIteration = 0;
        int m_nCentroidNotification = 10;
 
        public ContrastiveLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.CONTRASTIVE_LOSS;
 
            m_blobDiff = new Blob<T>(cuda, log, false);
            m_blobDiff.Name = m_param.name + " diff";
            m_blobDistSq = new Blob<T>(cuda, log, false);
            m_blobDistSq.Name = m_param.name + " distsq";
            m_blobDiffSq = new Blob<T>(cuda, log, false);
            m_blobDiffSq.Name = m_param.name + " diffsq";
            m_blobSummerVec = new Blob<T>(cuda, log, false);
            m_blobSummerVec.Name = m_param.name + " sum";
            m_blobSimilar = new Blob<T>(cuda, log, false);
            m_blobSimilar.Name = m_param.name + " similar";
            m_blobDistScale = new Blob<T>(cuda, log, false);
            m_blobDistScale.Name = m_param.name + " dist scale";
            m_blobPrimary = new Blob<T>(cuda, log, false);
            m_blobPrimary.Name = m_param.name + " primary";
        }
 
        protected override void dispose()
        {
            base.dispose();
 
            if (m_blobDiff != null)
            {
                m_blobDiff.Dispose();
                m_blobDiff = null;
            }
 
            if (m_blobDistSq != null)
            {
                m_blobDistSq.Dispose();
                m_blobDistSq = null;
            }
 
            if (m_blobDiffSq != null)
            {
                m_blobDiffSq.Dispose();
                m_blobDiffSq = null;
            }
 
            if (m_blobSummerVec != null)
            {
                m_blobSummerVec.Dispose();
                m_blobSummerVec = null;
            }
 
            if (m_blobSimilar != null)
            {
                m_blobSimilar.Dispose();
                m_blobSimilar = null;
            }
 
            if (m_blobDistScale != null)
            {
                m_blobDistScale.Dispose();
                m_blobDistScale = null;
            }
 
            if (m_blobPrimary != null)
            {
                m_blobPrimary.Dispose();
                m_blobPrimary = null;
            }
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return -1; }
        }
 
        public override int MinBottomBlobs
        {
            get { return 3; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 4; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return -1; }
        }
 
        public override int MinTopBlobs
        {
            get { return 1; }
        }
 
        public override int MaxTopBlobs
        {
            get { return 2; }
        }
 
        public override bool AllowForceBackward(int nBottomIdx)
        {
            if (nBottomIdx != 2)
                return true;
 
            return false;
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            m_nIteration = 0;
 
            m_log.CHECK_EQ(colBottom[0].channels, colBottom[1].channels, "the bottom[0] and bottom[1] should have equal channel values.");
            m_log.CHECK_EQ(1, colBottom[0].height, "The bottom[0] should have height = 1.");
            m_log.CHECK_EQ(1, colBottom[0].width, "The bottom[0] should have width = 1.");
            m_log.CHECK_EQ(1, colBottom[1].height, "The bottom[1] should have height = 1.");
            m_log.CHECK_EQ(1, colBottom[1].width, "The bottom[1] should have width = 1.");
            m_log.CHECK_GE(colBottom[2].channels, 1, "The bottom[2] should have channels >= 1.");
            m_log.CHECK_LE(colBottom[2].channels, 3, "The bottom[2] should have channels <= 3.");
            m_log.CHECK_EQ(1, colBottom[2].height, "The bottom[2] should have height = 1.");
            m_log.CHECK_EQ(1, colBottom[2].width, "The bottom[2] should have width = 1.");
            m_nCentroidNotification = 10;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
 
            m_blobDiff.Reshape(colBottom[0].num, colBottom[0].channels, 1, 1);
            m_blobDiffSq.Reshape(colBottom[0].num, colBottom[0].channels, 1, 1);
            m_blobDistSq.Reshape(colBottom[0].num, 1, 1, 1);
            // vector of ones used to sum along channels.
            m_blobSummerVec.Reshape(colBottom[0].channels, 1, 1, 1);
            m_blobSummerVec.SetData(1.0);
            m_blobPrimary.ReshapeLike(colBottom[0]);
 
            if (colTop.Count > 1 && m_param.contrastive_loss_param.output_matches)
            {
                if (m_rgMatches == null || m_rgMatches.Length != colBottom[0].num)
                    m_rgMatches = new T[colBottom[0].num];
 
                colTop[1].Reshape(colBottom[0].num, 1, 1, 1);
            }
 
            m_blobSimilar.Reshape(colBottom[0].num, 1, 1, 1);
            m_blobDistScale.Reshape(colBottom[0].num, 1, 1, 1);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nCount = colBottom[0].count();
 
            // Label data is in on of two forms:
            if (colBottom[2].channels > 1)
            {
                // channel > 1: the direct label values of each image packed into the data channels.
                m_cuda.channel_compare(colBottom[2].count(), colBottom[2].num, colBottom[2].channels, 1, colBottom[2].gpu_data, m_blobSimilar.mutable_gpu_data);
            }
            else
            {
                // channel = 1: the direct similarity where 1 = the same, and 0 = different.
                m_cuda.copy(colBottom[2].count(), colBottom[2].gpu_data, m_blobSimilar.mutable_gpu_data);
            }
 
            // When using centroid learning, the centroids from the DecodeLayer are only filled after they are fully
            // calculated - prior to full calculations, the centriods are set to 0.
            bool bCentroidLearningEnabled = false;
            bool bUseCentroidLearning = false;
            if (m_param.contrastive_loss_param.centroid_learning != ContrastiveLossParameter.CENTROID_LEARNING.NONE && colBottom.Count > 3)
            {
                bCentroidLearningEnabled = true;
                T fAsum = colBottom[3].asum_data();
                double dfAsum = convertD(fAsum);
                if (dfAsum > 0)
                    bUseCentroidLearning = true;
            }
 
            if (bUseCentroidLearning)
            {
                m_log.CHECK_EQ(colBottom.Count, 4, "When using centroid learning, a fourth bottom is required that contains the class centroids (calculated by the DecodeLayer).");
                m_log.CHECK_EQ(colBottom[3].channels, colBottom[0].channels, "Each centroid should have the same size as each encoding.");
                m_log.CHECK_EQ(colBottom[3].height, 1, "The centroids should have a height = 1.");
                m_log.CHECK_EQ(colBottom[3].width, 1, "The centroids should have a width = 1.");
                m_log.CHECK_EQ(colBottom[2].channels, 2, "The colBottom[2] must contain labels, not a similarity value - make sure the data layer has 'output_all_labels' = True.");
 
                // Load the target with the centroids to match the labels received in colBottom(2) - only use the first label of the two.
                int nEncodingDim = m_blobPrimary.count(1);
                int nLabelDim = colBottom[2].count(1);
 
                // Fill with centroids for each 'first' label.
                m_cuda.channel_fill(m_blobPrimary.count(), m_blobPrimary.num, nEncodingDim, 1, colBottom[3].gpu_data, nLabelDim, colBottom[2].gpu_data, m_blobPrimary.mutable_gpu_data);
 
                // If using centroid learning; for similar pairs, copy the centroids from colBottom[3], otherwise copy the colBottom[0] dissimilar encodings.
                if (m_param.contrastive_loss_param.centroid_learning == ContrastiveLossParameter.CENTROID_LEARNING.MATCHING)
                    m_cuda.copy(m_blobPrimary.count(), m_blobPrimary.num, m_blobPrimary.count(1), m_blobPrimary.gpu_data, colBottom[0].gpu_data, m_blobPrimary.mutable_gpu_data, m_blobSimilar.gpu_data); // centroid used for matching, use bottom[0] for all NON matching
 
                // If using centroid learning; for non-similar pairs, copy the centroids from colBottom[3], otherwise copy the colBottom[0] dissimilar encodings.
                else if (m_param.contrastive_loss_param.centroid_learning == ContrastiveLossParameter.CENTROID_LEARNING.NONMATCHING)
                    m_cuda.copy(m_blobPrimary.count(), m_blobPrimary.num, m_blobPrimary.count(1), m_blobPrimary.gpu_data, colBottom[0].gpu_data, m_blobPrimary.mutable_gpu_data, m_blobSimilar.gpu_data, true); // centroid used for NON matching, use bottom[0] for all matching.
 
                if (m_nCentroidNotification > 0)
                {
                    m_log.WriteLine("INFO: Centroid learning ON.");
                    m_nCentroidNotification--;
                }
            }
            else
            {
                // If not using centroid learning, just use the bottom[1] as is.
                m_cuda.copy(m_blobPrimary.count(), colBottom[0].gpu_data, m_blobPrimary.mutable_gpu_data);
            }
 
            bool bLegacyVersion = m_param.contrastive_loss_param.legacy_version;
            double dfMargin = m_param.contrastive_loss_param.margin;
            float[] rgSimPairs = Utility.ConvertVecF<T>(m_blobSimilar.update_cpu_data());
            float[] rgDist = null;
            double dfLoss = 0;
 
            if (m_param.contrastive_loss_param.distance_calculation == ContrastiveLossParameter.DISTANCE_CALCULATION.MANHATTAN)
            {
                // Manhattan Distance uses legacy calculation.
                bLegacyVersion = true;
 
                Blob<T> blobAbsDiff = m_blobDiffSq;
                Blob<T> blobDist = m_blobDistSq;
 
                m_cuda.sub(nCount,
                           m_blobPrimary.gpu_data,          // a
                           colBottom[1].gpu_data,           // b
                           m_blobDiff.mutable_gpu_data);    // a_i - b_i
 
                m_cuda.abs(nCount,
                           m_blobDiff.gpu_data,             //  a_i - b_i
                           blobAbsDiff.mutable_gpu_data);   // |a_i - b_i|
 
                m_cuda.gemv(false,
                           m_blobPrimary.num,
                           m_blobPrimary.channels,
                           1.0,
                           blobAbsDiff.gpu_data,           // |a_i - b_i|
                           m_blobSummerVec.gpu_data,
                           0.0,
                           blobDist.mutable_gpu_data);     // \Sum |a_i - b_i|
 
                if ((bUseCentroidLearning || !bCentroidLearningEnabled) && m_param.contrastive_loss_param.matching_distance_scale != 1.0)
                {
                    m_cuda.scale(m_blobDistScale.count(), m_param.contrastive_loss_param.matching_distance_scale - 1.0, m_blobSimilar.gpu_data, m_blobDistScale.mutable_gpu_data);
                    m_cuda.add_scalar(m_blobDistScale.count(), 1.0, m_blobDistScale.mutable_gpu_data);
                    m_cuda.mul(blobDist.count(), blobDist.gpu_data, m_blobDistScale.gpu_data, blobDist.mutable_gpu_data);
                }
 
                rgDist = Utility.ConvertVecF<T>(blobDist.update_cpu_data());
            }
            else // default = EUCLIDEAN
            {
                m_cuda.sub(nCount,
                           m_blobPrimary.gpu_data,          // a
                           colBottom[1].gpu_data,           // b
                           m_blobDiff.mutable_gpu_data);    // a_i - b_i
 
                m_cuda.powx(nCount,
                           m_blobDiff.mutable_gpu_data,     // a_i - b_i
                           2.0,
                           m_blobDiffSq.mutable_gpu_data);  // (a_i - b_i)^2
 
                m_cuda.gemv(false,
                           m_blobPrimary.num,
                           m_blobPrimary.channels,
                           1.0,
                           m_blobDiffSq.gpu_data,           // (a_i - b_i)^2
                           m_blobSummerVec.gpu_data,
                           0.0,
                           m_blobDistSq.mutable_gpu_data);  // \Sum (a_i - b_i)^2
 
                if ((bUseCentroidLearning || !bCentroidLearningEnabled) && m_param.contrastive_loss_param.matching_distance_scale != 1.0)
                {
                    m_cuda.scale(m_blobDistScale.count(), m_param.contrastive_loss_param.matching_distance_scale - 1.0, m_blobSimilar.gpu_data, m_blobDistScale.mutable_gpu_data);
                    m_cuda.add_scalar(m_blobDistScale.count(), 1.0, m_blobDistScale.mutable_gpu_data);
                    m_cuda.mul(m_blobDistSq.count(), m_blobDistSq.gpu_data, m_blobDistScale.gpu_data, m_blobDistSq.mutable_gpu_data);
                }
 
                rgDist = Utility.ConvertVecF<T>(m_blobDistSq.update_cpu_data());
            }
 
            for (int i = 0; i < colBottom[0].num; i++)
            {
                double dfDist = (bLegacyVersion) ? dfMargin - rgDist[i] : dfMargin - Math.Sqrt(rgDist[i]);
                bool bSimilar = (rgSimPairs[i] == 0) ? false : true;
 
                if (bSimilar)  // similar pairs
                {
                    if (m_rgMatches != null)
                    {
                        if (dfDist >= 0)
                            m_rgMatches[i] = m_tOne;
                        else
                            m_rgMatches[i] = m_tZero;
                    }
 
                    dfLoss += rgDist[i];
                }
                else // dissimilar pairs
                {
                    if (m_rgMatches != null)
                    {
                        if (dfDist >= 0)
                            m_rgMatches[i] = m_tZero;
                        else
                            m_rgMatches[i] = m_tOne;
                    }
 
                    dfDist = Math.Max(dfDist, 0);
 
                    if (bLegacyVersion)
                        dfLoss += dfDist;
                    else
                        dfLoss += dfDist * dfDist;
                }
            }
 
            dfLoss = dfLoss / (double)colBottom[0].num / 2.0;
            colTop[0].SetData(dfLoss, 0);
 
            if (colTop.Count > 1 && m_rgMatches != null)
                colTop[1].mutable_cpu_data = m_rgMatches;
 
            if (m_phase == Phase.TRAIN)
                m_nIteration++;
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0] && !rgbPropagateDown[1])
                return;
 
            double dfTopDiff = convertD(colTop[0].GetDiff(0)) / colBottom[0].num;
 
            m_log.CHECK_GT(m_blobSimilar.gpu_data, 0, "The similar data is not initialized - you must first run the forward pass under the Phase = TRAIN.");
 
            for (int i = 0; i < 2; i++)
            {
                if (rgbPropagateDown[i])
                {
                    int nCount = colBottom[0].count();
                    int nChannels = colBottom[0].channels;
                    double dfMargin = m_param.contrastive_loss_param.margin;
                    bool bLegacyVersion = m_param.contrastive_loss_param.legacy_version;
                    double dfSign = (i == 0) ? 1 : -1;
                    double dfAlpha = dfSign * dfTopDiff;
 
                    m_cuda.cll_bwd(nCount,
                                   nChannels,
                                   dfMargin,
                                   bLegacyVersion,
                                   dfAlpha,
                                   m_blobSimilar.gpu_data,  // pair similarity 0 or 1
                                   m_blobDiff.gpu_data,     // the cached eltwise difference between a and b
                                   m_blobDistSq.gpu_data,   // the cached square distance between a and b
                                   colBottom[i].mutable_gpu_diff);
                }
            }
        }
    }
}