Files
IronpenGPU/src/forward.jl
2023-07-30 10:12:24 +07:00

480 lines
19 KiB
Julia

module forward
# export
using Flux, CUDA
using GeneralUtils
using ..type, ..snnUtil
#------------------------------------------------------------------------------------------------100
""" kfn forward
input (row, col, batch)
"""
function (kfn::kfn_1)(input::AbstractArray)
kfn.timeStep .+= 1
#TODO time step forward
if view(kfn.learningStage, 1)[1] == 1
# reset learning params
# kfn.learningStage = [2]
end
# println(">>> input ", size(input))
# println(">>> zit ", size(kfn.zit))
# println(">>> lif_zit ", size(kfn.lif_zit))
# println(">>> lif_recSignal ", size(kfn.lif_recSignal))
# println(">>> lif_wRec ", size(kfn.lif_wRec))
# println(">>> lif_refractoryCounter ", size(kfn.lif_refractoryCounter))
# println(">>> lif_alpha ", size(kfn.lif_alpha))
# println(">>> lif_vt0 ", size(kfn.lif_vt0))
# println(">>> lif_vt0 sum ", sum(kfn.lif_vt0))
# pass input_data into input neuron.
GeneralUtils.cartesianAssign!(kfn.zit, input)
lifForward( kfn.zit,
kfn.lif_zit,
kfn.lif_wRec,
kfn.lif_vt0,
kfn.lif_vt1,
kfn.lif_vth,
kfn.lif_vRest,
kfn.lif_zt1,
kfn.lif_alpha,
kfn.lif_phi,
kfn.lif_epsilonRec,
kfn.lif_refractoryCounter,
kfn.lif_refractoryDuration,
kfn.lif_gammaPd,
kfn.lif_firingCounter,
kfn.lif_arrayProjection3DTo4D,
kfn.lif_recSignal,
kfn.lif_decayed_vt0,
kfn.lif_decayed_epsilonRec,
kfn.lif_vt1_diff_vth,
kfn.lif_vt1_diff_vth_div_vth,
kfn.lif_gammaPd_div_vth,
kfn.lif_phiActivation)
alifForward( kfn.zit,
kfn.alif_zit,
kfn.alif_wRec,
kfn.alif_vt0,
kfn.alif_vt1,
kfn.alif_vth,
kfn.alif_vRest,
kfn.alif_zt1,
kfn.alif_alpha,
kfn.alif_phi,
kfn.alif_epsilonRec,
kfn.alif_refractoryCounter,
kfn.alif_refractoryDuration,
kfn.alif_gammaPd,
kfn.alif_firingCounter,
kfn.alif_arrayProjection3DTo4D,
kfn.alif_recSignal,
kfn.alif_decayed_vt0,
kfn.alif_decayed_epsilonRec,
kfn.alif_vt1_diff_vth,
kfn.alif_vt1_diff_vth_div_vth,
kfn.alif_gammaPd_div_vth,
kfn.alif_phiActivation,
kfn.alif_epsilonRecA,
kfn.alif_avth,
kfn.alif_a,
kfn.alif_beta,
kfn.alif_rho,
kfn.alif_phi_x_epsilonRec,
kfn.alif_phi_x_beta,
kfn.alif_rho_diff_phi_x_beta,
kfn.alif_rho_div_phi_x_beta_x_epsilonRecA,
kfn.alif_beta_x_a)
# error("DEBUG -> kfn forward")
# update activation matrix by concatenate (input, lif_zt1, alif_zt1) to form activation matrix
_zit = cat(reshape(input, (size(input, 1), size(input, 2), 1, size(input, 3))),
reshape(kfn.lif_zt1, (size(input, 1), :, 1, size(input, 3))),
reshape(kfn.alif_zt1, (size(input, 1), :, 1, size(input, 3))), dims=2)
kfn.zit .= reshape(_zit, (size(input, 1), :, size(input, 3)))
# read out
onForward( kfn.zit,
kfn.on_zit,
kfn.on_wOut,
kfn.on_vt0,
kfn.on_vt1,
kfn.on_vth,
kfn.on_vRest,
kfn.on_zt1,
kfn.on_alpha,
kfn.on_phi,
kfn.on_epsilonRec,
kfn.on_refractoryCounter,
kfn.on_refractoryDuration,
kfn.on_gammaPd,
kfn.on_firingCounter,
kfn.on_arrayProjection3DTo4D,
kfn.on_recSignal,
kfn.on_decayed_vt0,
kfn.on_decayed_epsilonRec,
kfn.on_vt1_diff_vth,
kfn.on_vt1_diff_vth_div_vth,
kfn.on_gammaPd_div_vth,
kfn.on_phiActivation)
return reshape(kfn.on_zt1, (size(input, 1), :)),
kfn.zit
end
function lifForward(kfn_zit::Array{T},
zit::Array{T},
wRec::Array{T},
vt0::Array{T},
vt1::Array{T},
vth::Array{T},
vRest::Array{T},
zt1::Array{T},
alpha::Array{T},
phi::Array{T},
epsilonRec::Array{T},
refractoryCounter::Array{T},
refractoryDuration::Array{T},
gammaPd::Array{T},
firingCounter::Array{T},
arrayProjection3DTo4D::Array{T},
recSignal::Array{T},
decayed_vt0::Array{T},
decayed_epsilonRec::Array{T},
vt1_diff_vth::Array{T},
vt1_diff_vth_div_vth::Array{T},
gammaPd_div_vth::Array{T},
phiActivation::Array{T},
) where T<:Number
# project 3D kfn zit into 4D lif zit
zit .= reshape(kfn_zit,
(size(wRec, 1), size(wRec, 2), 1, size(wRec, 4))) .* arrayProjection3DTo4D
for j in 1:size(wRec, 4), i in 1:size(wRec, 3) # compute along neurons axis of every batch
if sum(@view(refractoryCounter[:,:,i,j])) > 0 # refractory period is active
@. @views refractoryCounter[:,:,i,j] -= 1
@. @views zt1[:,:,i,j] = 0
@. @views vt1[:,:,i,j] = alpha[:,:,i,j] * vt0[:,:,i,j]
@. @views phi[:,:,i,j] = 0
# compute epsilonRec
@. @views decayed_epsilonRec[:,:,i,j] = alpha[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views epsilonRec[:,:,i,j] = decayed_epsilonRec[:,:,i,j]
else # refractory period is inactive
@. @views recSignal[:,:,i,j] = zit[:,:,i,j] * wRec[:,:,i,j]
@. @views decayed_vt0[:,:,i,j] = alpha[:,:,i,j] * vt0[:,:,i,j]
@view(vt1[:,:,i,j]) .= @view(decayed_vt0[:,:,i,j]) .+ sum(@view(recSignal[:,:,i,j]))
if sum(@view(vt1[:,:,i,j])) > sum(@view(vth[:,:,i,j]))
@. @views zt1[:,:,i,j] = 1
@. @views refractoryCounter[:,:,i,j] = refractoryDuration[:,:,i,j]
@. @views firingCounter[:,:,i,j] += 1
@. @views vt1[:,:,i,j] = vRest[:,:,i,j]
else
@. @views zt1[:,:,i,j] = 0
end
# compute phi, there is a difference from alif formula
@. @views gammaPd_div_vth[:,:,i,j] = gammaPd[:,:,i,j] / vth[:,:,i,j]
@. @views vt1_diff_vth[:,:,i,j] = vt1[:,:,i,j] - vth[:,:,i,j]
@. @views vt1_diff_vth_div_vth[:,:,i,j] = vt1_diff_vth[:,:,i,j] / vth[:,:,i,j]
@view(phiActivation[:,:,i,j]) .= max(0, 1 - sum(@view(vt1_diff_vth_div_vth[:,:,i,j])))
@. @views phi[:,:,i,j] = gammaPd_div_vth[:,:,i,j] * phiActivation[:,:,i,j]
# compute epsilonRec
@. @views decayed_epsilonRec[:,:,i,j] = alpha[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views epsilonRec[:,:,i,j] = decayed_epsilonRec[:,:,i,j] + zit[:,:,i,j]
end
end
end
function alifForward(kfn_zit::Array{T},
zit::Array{T},
wRec::Array{T},
vt0::Array{T},
vt1::Array{T},
vth::Array{T},
vRest::Array{T},
zt1::Array{T},
alpha::Array{T},
phi::Array{T},
epsilonRec::Array{T},
refractoryCounter::Array{T},
refractoryDuration::Array{T},
gammaPd::Array{T},
firingCounter::Array{T},
arrayProjection3DTo4D::Array{T},
recSignal::Array{T},
decayed_vt0::Array{T},
decayed_epsilonRec::Array{T},
vt1_diff_vth::Array{T},
vt1_diff_vth_div_vth::Array{T},
gammaPd_div_vth::Array{T},
phiActivation::Array{T},
epsilonRecA::Array{T},
avth::Array{T},
a::Array{T},
beta::Array{T},
rho::Array{T},
phi_x_epsilonRec::Array{T},
phi_x_beta::Array{T},
rho_diff_phi_x_beta::Array{T},
rho_div_phi_x_beta_x_epsilonRecA::Array{T},
beta_x_a::Array{T},
) where T<:Number
# project 3D kfn zit into 4D lif zit
zit .= reshape(kfn_zit,
(size(wRec, 1), size(wRec, 2), 1, size(wRec, 4))) .* arrayProjection3DTo4D
for j in 1:size(wRec, 4), i in 1:size(wRec, 3) # compute along neurons axis of every batch
if sum(@view(refractoryCounter[:,:,i,j])) > 0 # refractory period is active
@. @views refractoryCounter[:,:,i,j] -= 1
@. @views zt1[:,:,i,j] = 0
@. @views vt1[:,:,i,j] = alpha[:,:,i,j] * vt0[:,:,i,j]
@. @views phi[:,:,i,j] = 0
@. @views a[:,:,i,j] = rho[:,:,i,j] * a[:,:,i,j]
# compute epsilonRec
@. @views decayed_epsilonRec[:,:,i,j] = alpha[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views epsilonRec[:,:,i,j] = decayed_epsilonRec[:,:,i,j]
# compute epsilonRecA
@. @views phi_x_epsilonRec[:,:,i,j] = phi[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views phi_x_beta[:,:,i,j] = phi[:,:,i,j] * beta[:,:,i,j]
@. @views rho_diff_phi_x_beta[:,:,i,j] = rho[:,:,i,j] - phi_x_beta[:,:,i,j]
@. @views rho_div_phi_x_beta_x_epsilonRecA[:,:,i,j] = rho_diff_phi_x_beta[:,:,i,j] * epsilonRecA[:,:,i,j]
@. @views epsilonRecA[:,:,i,j] = phi_x_epsilonRec[:,:,i,j] + rho_div_phi_x_beta_x_epsilonRecA[:,:,i,j]
# compute avth
@. @views beta_x_a[:,:,i,j] = beta[:,:,i,j] * a[:,:,i,j]
@. @views avth[:,:,i,j] = vth[:,:,i,j] + beta_x_a[:,:,i,j]
else # refractory period is inactive
@. @views recSignal[:,:,i,j] = zit[:,:,i,j] * wRec[:,:,i,j]
@. @views decayed_vt0[:,:,i,j] = alpha[:,:,i,j] * vt0[:,:,i,j]
@view(vt1[:,:,i,j]) .= @view(decayed_vt0[:,:,i,j]) .+ sum(@view(recSignal[:,:,i,j]))
# compute avth
@. @views beta_x_a[:,:,i,j] = beta[:,:,i,j] * a[:,:,i,j]
@. @views avth[:,:,i,j] = vth[:,:,i,j] + beta_x_a[:,:,i,j]
if sum(@view(vt1[:,:,i,j])) > sum(@view(avth[:,:,i,j]))
@. @views zt1[:,:,i,j] = 1
@. @views refractoryCounter[:,:,i,j] = refractoryDuration[:,:,i,j]
@. @views firingCounter[:,:,i,j] += 1
@. @views vt1[:,:,i,j] = vRest[:,:,i,j]
@. @views a[:,:,i,j] = rho[:,:,i,j] * a[:,:,i,j]
@. @views a[:,:,i,j] = a[:,:,i,j] += 1
else
@. @views zt1[:,:,i,j] = 0
@. @views a[:,:,i,j] = rho[:,:,i,j] * a[:,:,i,j]
end
# compute phi, there is a difference from alif formula
@. @views gammaPd_div_vth[:,:,i,j] = gammaPd[:,:,i,j] / vth[:,:,i,j]
@. @views vt1_diff_vth[:,:,i,j] = vt1[:,:,i,j] - vth[:,:,i,j]
@. @views vt1_diff_vth_div_vth[:,:,i,j] = vt1_diff_vth[:,:,i,j] / vth[:,:,i,j]
@view(phiActivation[:,:,i,j]) .= max(0, 1 - sum(@view(vt1_diff_vth_div_vth[:,:,i,j])))
@. @views phi[:,:,i,j] = gammaPd_div_vth[:,:,i,j] * phiActivation[:,:,i,j]
# compute epsilonRec
@. @views decayed_epsilonRec[:,:,i,j] = alpha[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views epsilonRec[:,:,i,j] = decayed_epsilonRec[:,:,i,j] + zit[:,:,i,j]
# compute epsilonRecA
@. @views phi_x_epsilonRec[:,:,i,j] = phi[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views phi_x_beta[:,:,i,j] = phi[:,:,i,j] * beta[:,:,i,j]
@. @views rho_diff_phi_x_beta[:,:,i,j] = rho[:,:,i,j] - phi_x_beta[:,:,i,j]
@. @views rho_div_phi_x_beta_x_epsilonRecA[:,:,i,j] = rho_diff_phi_x_beta[:,:,i,j] * epsilonRecA[:,:,i,j]
@. @views epsilonRecA[:,:,i,j] = phi_x_epsilonRec[:,:,i,j] + rho_div_phi_x_beta_x_epsilonRecA[:,:,i,j]
end
end
end
function onForward(kfn_zit::Array{T},
zit::Array{T},
wOut::Array{T},
vt0::Array{T},
vt1::Array{T},
vth::Array{T},
vRest::Array{T},
zt1::Array{T},
alpha::Array{T},
phi::Array{T},
epsilonRec::Array{T},
refractoryCounter::Array{T},
refractoryDuration::Array{T},
gammaPd::Array{T},
firingCounter::Array{T},
arrayProjection3DTo4D::Array{T},
recSignal::Array{T},
decayed_vt0::Array{T},
decayed_epsilonRec::Array{T},
vt1_diff_vth::Array{T},
vt1_diff_vth_div_vth::Array{T},
gammaPd_div_vth::Array{T},
phiActivation::Array{T},
) where T<:Number
# project 3D kfn zit into 4D lif zit
zit .= reshape(kfn_zit,
(size(wOut, 1), size(wOut, 2), 1, size(wOut, 4))) .* arrayProjection3DTo4D
for j in 1:size(wOut, 4), i in 1:size(wOut, 3) # compute along neurons axis of every batch
if sum(@view(refractoryCounter[:,:,i,j])) > 0 # refractory period is active
@. @views refractoryCounter[:,:,i,j] -= 1
@. @views zt1[:,:,i,j] = 0
@. @views vt1[:,:,i,j] = alpha[:,:,i,j] * vt0[:,:,i,j]
@. @views phi[:,:,i,j] = 0
# compute epsilonRec
@. @views decayed_epsilonRec[:,:,i,j] = alpha[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views epsilonRec[:,:,i,j] = decayed_epsilonRec[:,:,i,j]
else # refractory period is inactive
@. @views recSignal[:,:,i,j] = zit[:,:,i,j] * wOut[:,:,i,j]
@. @views decayed_vt0[:,:,i,j] = alpha[:,:,i,j] * vt0[:,:,i,j]
@view(vt1[:,:,i,j]) .= @view(decayed_vt0[:,:,i,j]) .+ sum(@view(recSignal[:,:,i,j]))
if sum(@view(vt1[:,:,i,j])) > sum(@view(vth[:,:,i,j]))
@. @views zt1[:,:,i,j] = 1
@. @views refractoryCounter[:,:,i,j] = refractoryDuration[:,:,i,j]
@. @views firingCounter[:,:,i,j] += 1
@. @views vt1[:,:,i,j] = vRest[:,:,i,j]
else
@. @views zt1[:,:,i,j] = 0
end
# compute phi, there is a difference from alif formula
@. @views gammaPd_div_vth[:,:,i,j] = gammaPd[:,:,i,j] / vth[:,:,i,j]
@. @views vt1_diff_vth[:,:,i,j] = vt1[:,:,i,j] - vth[:,:,i,j]
@. @views vt1_diff_vth_div_vth[:,:,i,j] = vt1_diff_vth[:,:,i,j] / vth[:,:,i,j]
@view(phiActivation[:,:,i,j]) .= max(0, 1 - sum(@view(vt1_diff_vth_div_vth[:,:,i,j])))
@. @views phi[:,:,i,j] = gammaPd_div_vth[:,:,i,j] * phiActivation[:,:,i,j]
# compute epsilonRec
@. @views decayed_epsilonRec[:,:,i,j] = alpha[:,:,i,j] * epsilonRec[:,:,i,j]
@. @views epsilonRec[:,:,i,j] = decayed_epsilonRec[:,:,i,j] + zit[:,:,i,j]
end
end
end
# function onForward(kfn_zit,
# zit,
# wOut,
# vt0,
# vt1,
# vth,
# vRest,
# zt1,
# alpha,
# phi,
# epsilonRec,
# refractoryCounter,
# refractoryDuration,
# gammaPd,
# firingCounter)
# d1, d2, d3, d4 = size(wOut)
# zit .= reshape(kfn_zit, (d1, d2, 1, d4)) .* ones(size(wOut)...) # project zit into zit
# for j in 1:d4, i in 1:d3 # compute along neurons axis of every batch
# if view(refractoryCounter, :, :, i, j)[1] > 0 # neuron is inactive (in refractory period)
# view(refractoryCounter, :, :, i, j)[1] -= 1
# view(zt1, :, :, i, j)[1] = 0
# view(vt1, :, :, i, j)[1] =
# view(alpha, :, :, i, j)[1] * view(vt0, :, :, i, j)[1]
# view(phi, :, :, i, j)[1] = 0.0
# view(epsilonRec, :, :, i, j) .= view(alpha, :, :, i, j)[1] .*
# view(epsilonRec, :, :, i, j)
# else # neuron is active
# view(vt1, :, :, i, j)[1] =
# (view(alpha, :, :, i, j)[1] * view(vt0,:, :, i, j)[1]) +
# sum(view(zit, :, :, i, j) .* view(wOut, :, :, i, j))
# if view(vt1, :, :, i, j)[1] > view(vth, :, :, i, j)[1]
# view(zt1, :, :, i, j)[1] = 1
# view(refractoryCounter, :, :, i, j)[1] =
# view(refractoryDuration, :, :, i, j)[1]
# view(firingCounter, :, :, i, j)[1] += 1
# view(vt1, :, :, i, j)[1] = view(vRest, :, :, i, j)[1]
# else
# view(zt1, :, :, i, j)[1] = 0
# end
# # there is a difference from alif formula
# view(phi, :, :, i, j)[1] =
# (view(gammaPd, :, :, i, j)[1] / view(vth, :, :, i, j)[1]) *
# max(0, 1 - ((view(vt1, :, :, i, j)[1] - view(vth, :, :, i, j)[1]) /
# view(vth, :, :, i, j)[1]))
# view(epsilonRec, :, :, i, j) .=
# (view(alpha, :, :, i, j)[1] .* view(epsilonRec, :, :, i, j)) +
# view(zit, :, :, i, j)
# end
# end
# end
end # module