pyloss.py

# -*- coding: utf-8 -*-
"""
Created on Mon Apr 18 13:30:19 2016

@author: joe
"""

import sys
sys.path.insert(0, '/home/joe/github/caffe-4-17/caffe/python/')

import caffe
import numpy as np

class EuclideanLossLayer(caffe.Layer):
    
    def setup(self, bottom, top):
        print('setup()')
        if len(bottom) != 2:
            raise Exception("Need two inputs to compute distance.")
            
    def reshape(self, bottom, top):
        print('reshape()')
        if bottom[0].count != bottom[1].count:
            raise Exception("Inputs must have the same dimension.")
        # difference is shape of inputs
        self.diff = np.zeros_like(bottom[0].data, dtype=np.float32)
        # loss output is scalar
        top[0].reshape(1)
        
    def forward(self, bottom, top):
        print('forward()')
        self.diff[...] = bottom[0].data - bottom[1].data
        top[0].data[...] = np.sum(self.diff**2) / bottom[0].num / 2.
        
    def backward(self, top, propagate_down, bottom):
        print('backward()')
        for i in range(2):
            if not propagate_down[i]:
                continue
            if i == 0:
                sign = 1
            else:
                sign = -1
            bottom[i].diff[...] = sign * self.diff / bottom[i].num