ton/IronpenGPU
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dev

Browse Source
This commit is contained in:
ton
2023-09-19 20:45:17 +07:00
parent 32de83caf3
commit b073649129
5 changed files with 136 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/forward.jl
View File

@@ -20,7 +20,7 @@ function (kfn::kfn_1)(input::AbstractArray)
kfn.timeStep .= 1
# reset learning params
kfn.zitCumulative = kfn.zitCumulative[:,:,1,:]
kfn.zitCumulative = (kfn.zitCumulative[:,:,1] .= 0)
kfn.lif_vt .= 0
kfn.lif_wRecChange .= 0
@@ -44,6 +44,7 @@ function (kfn::kfn_1)(input::AbstractArray)
kfn.on_epsilonRec .= 0
kfn.on_wOutChange .= 0
kfn.on_refractoryCounter .= 0
kfn.on_synapticActivityCounter .= 0
kfn.learningStage = [2]
end
@@ -153,6 +154,7 @@ function (kfn::kfn_1)(input::AbstractArray)
kfn.on_gammaPd,
kfn.on_firingCounter,
kfn.on_recSignal,
kfn.on_synapticActivityCounter,
)
# get on_zt4d to on_zt
kfn.on_zt .= reduce(max, kfn.on_zt4d, dims=(1,2))
@@ -557,6 +559,7 @@ function onForward( zit::CuArray,
gammaPd::CuArray,
firingCounter::CuArray,
recSignal::CuArray,
synapticActivityCounter::CuArray,
)
kernel = @cuda launch=false onForward( zit,
@@ -573,6 +576,7 @@ function onForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
synapticActivityCounter,
GeneralUtils.linear_to_cartesian,
)
config = launch_configuration(kernel.fun)
@@ -602,6 +606,7 @@ function onForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
synapticActivityCounter,
GeneralUtils.linear_to_cartesian; threads, blocks)
end
end
@@ -621,6 +626,7 @@ function onForward( zit,
gammaPd,
firingCounter,
recSignal,
synapticActivityCounter,
linear_to_cartesian,
)
i = (blockIdx().x - 1) * blockDim().x + threadIdx().x # gpu threads index
@@ -661,6 +667,8 @@ function onForward( zit,
# compute epsilonRec
epsilonRec[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4]) +
(zit[i1,i2,i3,i4] * !iszero(wOut[i1,i2,i3,i4]))
synapticActivityCounter[i1,i2,i3,i4] += zit[i1,i2,i3,i4] * !iszero(wOut[i1,i2,i3,i4])
end
end
return nothing