blob: 3a23fdb6447c94c06abac8f6e25085897d4fdd3b [file] [log] [blame]
#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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.
#
#-------------------------------------------------------------
/*
* 2D Softmax classifier layer.
*/
source("scripts/nn/util.dml") as util
source("nn/layers/softmax.dml") as softmax
forward = function(matrix[double] scores, int C)
return (matrix[double] probs) {
/*
* Computes the forward pass for a softmax2d classifier. The input
* has four dimensions (N, C, Hin, Win), that means it has N
* 2d-examples with a shape (Hin, Win), each pixel in the 2d
* example has C values that are interpreted as unnormalized,
* log-probabilities for each of C classes. The softmax function
* transforms these values to normalized probabilities across the C
* classes, for every example.
*
* This can be interpreted as a generalization of the sigmoid
* function to multiple classes.
*
* `probs_ijk = e^scores_ijk / sum(e^scores_ijk)`
*
* In these equations, `probs_ijk` is the C-dimensional vector of the
* normalized probabilities for the pixel `j, k` in the example `i`
*
* Inputs:
* - scores: Inputs, of shape (N, C*Hin*Win).
* - C: Number of input channels (dimensionality of input depth).
*
* Outputs:
* - probs: Outputs, of shape (N, C*Hin*Win).
*/
# For numerical stability, we subtract the max score of an example from all scores for that
# example. This is equivalent to the original formulation:
# e^scores_ijk / sum(e^scores_ijk) == C*e^scores_ijk / C*sum(e^scores_ijk)
# == e^(scores_ijk+log(C)) / sum(e^(scores_ijk+log(C))
# set log(C) = -max(scores_ijk):
# == e^(scores_ijk-max(scores_ijk)) / sum(e^(scores_ijk-max(scores_ijk))
N = nrow(scores)
#Transpose the matrix from (N, C*H*W) to (N*H*W, C)
scores_C_NHW = util::transpose_NCHW_to_CNHW(scores, C)
scores_NHW_C = t(scores_C_NHW)
probs_NHW_C = softmax::forward(scores_NHW_C)
#Transpose the matrix from (N*H*W, C) to (N, C*H*W)
probs_C_NHW = t(probs_NHW_C)
probs = util::transpose_NCHW_to_CNHW(probs_C_NHW, N)
}
backward = function(matrix[double] dprobs, matrix[double] scores, int C)
return (matrix[double] dscores) {
/*
* Computes the backward pass for a softmax2d classifier.
*
* Note that dscores_ij has multiple source branches:
*
* ```
* dprobs_ij/dscores_ij = probs_ij * (1 - probs_ij)
* dprobs_ik/dscores_ij = -probs_ik * probs_ij, for all k != j
*
* dloss/dscores_ij =
* (dloss/dprobs_ij * dprobs_ij/dscores_ij)
* + sum_{k!=j}(dloss/dprobs_ik * dprobs_ik/dscores_ij)
* ```
*
* Inputs:
* - dprobs: Gradient wrt `probs` from upstream, of shape (N, C*Hin*Win).
* - scores: Inputs, of shape (N, C*Hin*Win).
* - C: Number of input channels (dimensionality of input depth).
*
* Outputs:
* - dscores: Gradient wrt `scores`, of shape (N, C*Win*Hin).
*/
N = nrow(scores)
#Transpose the matrix from (N, C*H*W) to (N*H*W, C)
dprobs_C_NHW = util::transpose_NCHW_to_CNHW(dprobs, C)
dprobs_NHW_C = t(dprobs_C_NHW)
#Transpose the matrix from (N, C*H*W) to (N*H*W, C)
scores_C_NHW = util::transpose_NCHW_to_CNHW(scores, C)
scores_NHW_C = t(scores_C_NHW)
dscores_NHW_C = softmax::backward(dprobs_NHW_C, scores_NHW_C)
#Transpose the matrix from (N*H*W, C) to (N, C*H*W)
dscores_C_NHW = t(dscores_NHW_C)
dscores = util::transpose_NCHW_to_CNHW(dscores_C_NHW, N)
}