As I understand, you want to change the validInputSize to “[x x numClasses]”, which implies a single prediction will be of size [x x 1]. Hence, you also need to reshape your classWeights (which is W in your code) to [x x numClasses].
Modified your code below:
classdef weightedClassificationLayer < nnet.layer.ClassificationLayer
properties
% Row vector of weights corresponding to the classes in the
% training data.
ClassWeights
end
methods
function layer = weightedClassificationLayer(classWeights, name)
% layer = weightedClassificationLayer(classWeights) creates a
% weighted cross entropy loss layer. classWeights is a row
% vector of weights corresponding to the classes in the order
% that they appear in the training data.
% layer = weightedClassificationLayer(classWeights, name)
% additionally specifies the layer name
% Set class weights.
layer.ClassWeights = classWeights;
% Set layer name.
if nargin == 2
layer.Name = name;
end
% Set layer description
layer.Description = 'Weighted cross entropy';
end
function loss = forwardLoss(layer, Y, T)
% loss = forwardLoss(layer, Y, T) returns the weighted cross
% entropy loss between the predictions Y and the training
% targets T.
N = size(Y,4);
Y = squeeze(Y);
T = squeeze(T);
W = layer.ClassWeights;
%% Modified %%
n = length(T(:))/2;
W = repelem(W, 1, [n n]);
W = reshape(W, size(T));
prod = W.*(T.*log(Y));
loss = -sum(prod(:))/N;
%% Modified %%
end
function dLdY = backwardLoss(layer, Y, T)
% dLdX = backwardLoss(layer, Y, T) returns the derivatives of
% the weighted cross entropy loss with respect to the
% predictions Y.
[H,Wi,K,N] = size(Y);
Y = squeeze(Y);
T = squeeze(T);
W = layer.ClassWeights;
%% Modified %%
n = length(T(:))/2;
W = repelem(W, 1, [n n]);
W = reshape(W, size(T));
%% Modified %%
dLdY = -(W.*T./Y)/N;
dLdY = reshape(dLdY,[H Wi K N]);
end
end
end
