% Version 1.000
%
% Code provided by Ruslan Salakhutdinov
%
% Permission is granted for anyone to copy, use, modify, or distribute this
% program and accompanying programs and documents for any purpose, provided
% this copyright notice is retained and prominently displayed, along with
% a note saying that the original programs are available from our
% web page.
% The programs and documents are distributed without any warranty, express or
% implied. As the programs were written for research purposes only, they have
% not been tested to the degree that would be advisable in any important
% application. All use of these programs is entirely at the user's own risk.

% This program trains Restricted Boltzmann Machine in which
% visible, binary, stochastic pixels are connected to
% hidden, binary, stochastic feature detectors using symmetrically
% weighted connections. Learning is done with 1-step Contrastive Divergence.
% The program assumes that the following variables are set externally:
% maxepoch -- maximum number of epochs
% numhid -- number of hidden units
% batchdata -- the data that is divided into batches (numcases numdims numbatches)
% restart -- set to 1 if learning starts from beginning

function rbm(maxepoch, numhid, batchdata, restart)
if restart == 1,
%restart=0;

epsilonw = 0.05; % Learning rate for weights
epsilonvb = 0.05; % Learning rate for biases of visible units
epsilonhb = 0.05; % Learning rate for biases of hidden units

%CD=1;
weightcost = 0.001;
initialmomentum = 0.5;
finalmomentum = 0.9;

[batchSize, nPixels, nBatches] = size(batchdata);
epoch=1;

% Initializing symmetric weights and biases.
vishid = 0.001*randn(nPixels, numhid);
hidbiases = zeros(1,numhid);
visbiases = zeros(1,nPixels);

% expVisAct = zeros(numcases,numhid);
% expHidAct = zeros(numcases,numhid);
% expVisAct_neg = zeros(numcases,numhid);
% expHidAct_neg = zeros(numcases,numhid);
% avSynapseStrength_clampingV = zeros(numdims,numhid);
% avSynapseStrength_threshold = zeros(numdims,numhid);

vishidinc = zeros(nPixels,numhid);
hidbiasinc = zeros(1,numhid);
visbiasinc = zeros(1,nPixels);
end

for epoch = epoch:maxepoch
fprintf(1,'epoch %d\r',epoch);
errsum=0;
for batch = 1:nBatches,
if mod(batch,10) == 9,
fprintf(1,'.');
end
visbias = repmat( visbiases,batchSize,1);
hidbias = repmat(2*hidbiases,batchSize,1);
%%%%%%%%% START POSITIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
data = batchdata(:,:,batch);
% By considering data (pixel value in [0.0, 1.0]) as probability,
% convert to 0 or 1 accordingly to get VIS outputs.
sampleFromVis = data > rand(batchSize,nPixels);
% clear data

% Clamping VIS, get HID layer activations, i.e. POSTERIOR dist. for HID
hidDist = sigma( sampleFromVis*(2*vishid) + hidbias ); % Eq. 7
% ^ pos for POSitive not POSterior (although it IS the posterior)

% 784 vis/inputs going to 10 hiddens = a 784x10 v_i h_j matrix.
% Sum over all batch items.
avSynapseStrength_clampingV = sampleFromVis' * hidDist / batchSize; % 784x10

% Get expected activations for each visible and hidden neuron
expVisAct = sum(sampleFromVis) / batchSize; % <v_i h_j>_data in Eq. 5
expHidAct = sum(hidDist) / batchSize;

%%%%%%%%% END OF POSITIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%% START NEGATIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Sample from our posterior dist. for HID
sampleFromHid = hidDist > rand(batchSize,numhid);
% Clamping HID to this sample, get a POSTERIOR dist. for VIS ...
visRecoProbs = sigma( sampleFromHid*vishid' + visbias ); % Eq. 8
% ... from which we sample a 'Reconstruction' vis/data-vector ...
visReco = visRecoProbs > rand(batchSize,nPixels);
% ... and throw it back thru hidden units to get a 'neg' HID dist.
newHidDist = sigma( visReco*(2*vishid) + hidbias );

avSynapseStrength_threshold = visReco'*newHidDist / batchSize;

expVisAct_neg = sum(visReco) / batchSize;
expHidAct_neg = sum(newHidDist) / batchSize;

% clear sampleFromHid visRecoProbs hidDist newHidDist

%%%%%%%%% END OF NEGATIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% keep a tally of total divergence of reconstructions from
% original sampled vis-vector.
err = sum(sum( (sampleFromVis-visReco).^2 ));
errsum = err + errsum;

% clear hidDist newHidDist sampleFromVis

if epoch>5,
momentum=finalmomentum;
else
momentum=initialmomentum;
end;

%%%%%%%%% UPDATE WEIGHTS AND BIASES %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
avSynapseExcitation = avSynapseStrength_clampingV - avSynapseStrength_threshold

vishidinc = momentum*vishidinc + epsilonw * ( avSynapseExcitation - weightcost*vishid);
visbiasinc = momentum*visbiasinc + epsilonvb * (expVisAct - expVisAct_neg);
hidbiasinc = momentum*hidbiasinc + epsilonhb * (expHidAct - expHidAct_neg);

vishid = vishid + vishidinc;
visbiases = visbiases + visbiasinc;
hidbiases = hidbiases + hidbiasinc;
%%%%%%%%%%%%%%%% END OF UPDATES %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if rem(batch,600)==0
figure(1);
dispims(visReco',28,28);
drawnow
end
end
fprintf(1, ' error %6.1f \n', errsum);

end

save data/out/1_fullmnistvh.mat vishid visbiases hidbiases epoch

end