Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

🔨 refactor(layer_seq): SelectSeq #82

Merged
merged 3 commits into from
Apr 3, 2023
Merged

🔨 refactor(layer_seq): SelectSeq #82

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
adfc13d
SelectSeq
jean-francoisreboud Mar 31, 2023
f529186
Update changelog [skip ci]
jean-francoisreboud Mar 31, 2023
9353710
Merge branch 'release_4'
jean-francoisreboud Apr 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.

## [unreleased]

🔨 **layer_seq**: SelectSeq ([#82](https://github.com/owkin/GrAIdient/pull/82))\
🚀 **examples**: AutoEncoder models ([#79](https://github.com/owkin/GrAIdient/pull/79))\
🪜 **layer_2d**: VQ2D ([#81](https://github.com/owkin/GrAIdient/pull/81))\
🚀 **layer_seq**: factorize by nbHeads ([#78](https://github.com/owkin/GrAIdient/pull/78))\
Expand Down
122 changes: 19 additions & 103 deletions ...GrAIdient/LayerSeq/SelectNeuronsSeq.swift → Sources/GrAIdient/LayerSeq/SelectSeq.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// SelectNeuronsSeq.swift
// SelectSeq.swift
// GrAIdient
//
// Created by Aurélien PEDEN on 17/03/2023.
Expand All @@ -10,32 +10,14 @@
///
/// This layer selects one element of a sequence and transforms a LayerSeq into a Layer1D.
///
public class SelectNeuronsSeq: Layer1D
public class SelectSeq: Layer1D
{
/// Sequence of the selected neurons.
let _targetSeq: Int

/// List of neurons to select.
let _neurons: [Int]
/// List of coefficients to scale each selected neuron.
let _coeffs: [Double]

///
/// Indices of selected neurons.
/// Shape ~ (nbNeurons,).
///
var _neuronsBuffer: MetalPrivateBuffer<UInt32>! = nil
///
/// Coefficients of selected neurons.
/// Shape ~ (nbNeurons,).
///
var _coeffsBuffer: MetalPrivateBuffer<Float>! = nil

private enum Keys: String, CodingKey
{
case targetSeq
case neurons
case coeffs
}

///
Expand All @@ -44,21 +26,15 @@ public class SelectNeuronsSeq: Layer1D
/// - Parameters:
/// - layerPrev: Previous layer that has been queued to the model.
/// - targetSeq: Sequence of the selected neurons.
/// - neurons: The list of neurons to select.
/// - coeffs: The list of coefficients to scale each selected neuron.
/// - params: Contextual parameters linking to the model.
///
public init(layerPrev: LayerSeq,
targetSeq: Int,
neurons: [Int], coeffs: [Double],
params: GrAI.Model.Params)
{
_targetSeq = targetSeq
_neurons = neurons
_coeffs = coeffs

super.init(layerPrev: layerPrev,
nbNeurons: _neurons.count,
nbNeurons: layerPrev.nbNeurons,
params: params)
}

Expand All @@ -74,8 +50,6 @@ public class SelectNeuronsSeq: Layer1D
{
let values = try decoder.container(keyedBy: Keys.self)
_targetSeq = try values.decode(Int.self, forKey: Keys.targetSeq)
_neurons = try values.decode([Int].self, forKey: .neurons)
_coeffs = try values.decode([Double].self, forKey: .coeffs)
try super.init(from: decoder)
}

Expand All @@ -94,8 +68,6 @@ public class SelectNeuronsSeq: Layer1D
{
var container = encoder.container(keyedBy: Keys.self)
try container.encode(_targetSeq, forKey: Keys.targetSeq)
try container.encode(_neurons, forKey: Keys.neurons)
try container.encode(_coeffs, forKey: Keys.coeffs)
try super.encode(to: encoder)
}

Expand All @@ -121,60 +93,14 @@ public class SelectNeuronsSeq: Layer1D
let params = GrAI.Model.Params(context: context)
params.context.curID = id

let layer = SelectNeuronsSeq(
let layer = SelectSeq(
layerPrev: layerPrev,
targetSeq: _targetSeq,
neurons: _neurons,
coeffs: _coeffs,
params: params
)
return layer
}

///
/// Clean state resources in the GPU execution context.
///
/// We clean the neurons' state (forward and backward).
///
public override func resetKernelGPU()
{
super.resetKernelGPU()
_neuronsBuffer = nil
_coeffsBuffer = nil
}

///
/// Initialize state resources in the GPU execution context.
///
/// We initialize the neurons' forward state.
///
public override func checkStateForwardGPU(batchSize: Int) throws
{
try super.checkStateForwardGPU(batchSize: batchSize)

if _neuronsBuffer == nil
{
_neuronsBuffer = MetalPrivateBuffer<UInt32>(
nbNeurons, deviceID: deviceID
)
_coeffsBuffer = MetalPrivateBuffer<Float>(
nbNeurons, deviceID: deviceID
)

let neuronsPtr = _neuronsBuffer.shared.buffer
let coeffsPtr = _coeffsBuffer.shared.buffer

for (num, neuron) in _neurons.enumerated()
{
neuronsPtr[num] = UInt32(neuron)
coeffsPtr[num] = Float(_coeffs[num])
}

MetalKernel.get.upload([_neuronsBuffer])
MetalKernel.get.upload([_coeffsBuffer])
}
}

///
/// Apply the forward pass of the Gradient Checking in CPU execution context.
///
Expand All @@ -198,9 +124,9 @@ public class SelectNeuronsSeq: Layer1D
{
for depth in 0..<nbNeurons
{
let outPrev = neuronsPrev.get(_targetSeq, _neurons[depth])!
let outPrev = neuronsPrev.get(_targetSeq, depth)!
neurons.get(depth)!.gc[batch][elem].out =
_coeffs[depth] * outPrev.gc[batch][elem].out
outPrev.gc[batch][elem].out
}
}}
}
Expand Down Expand Up @@ -233,8 +159,7 @@ public class SelectNeuronsSeq: Layer1D
for depth in 0..<nbNeurons
{
neurons.get(depth)!.v[elem].out =
_coeffs[depth] * neuronsPrev
.get(_targetSeq, _neurons[depth])!.v[elem].out
neuronsPrev.get(_targetSeq, depth)!.v[elem].out
}
}
}
Expand All @@ -253,22 +178,18 @@ public class SelectNeuronsSeq: Layer1D

let pTargetSeq: [UInt32] = [UInt32(_targetSeq)]
let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
let pNbNeuronsPrev: [UInt32] = [UInt32(layerPrev.nbNeurons)]
let pNbBatch: [UInt32] = [UInt32(batchSize)]
let pSequence: [UInt32] = [UInt32(layerPrev.sequence)]

let command = MetalKernel.get.createCommand(
"selectNeuronsSeqForward", deviceID: deviceID
"selectSeqForward", deviceID: deviceID
)
command.setBuffer(layerPrev.outs.metal, atIndex: 0)
command.setBytes(pNbNeurons, atIndex: 1)
command.setBytes(pNbNeuronsPrev, atIndex: 2)
command.setBytes(pTargetSeq, atIndex: 3)
command.setBuffer(_neuronsBuffer.metal, atIndex: 4)
command.setBuffer(_coeffsBuffer.metal, atIndex: 5)
command.setBytes(pNbBatch, atIndex: 6)
command.setBytes(pSequence, atIndex: 7)
command.setBuffer(outs.metal, atIndex: 8)
command.setBytes(pTargetSeq, atIndex: 2)
command.setBytes(pNbBatch, atIndex: 3)
command.setBytes(pSequence, atIndex: 4)
command.setBuffer(outs.metal, atIndex: 5)

command.dispatchThreads(
width: nbNeurons,
Expand Down Expand Up @@ -300,9 +221,8 @@ public class SelectNeuronsSeq: Layer1D
{
for depth in 0..<nbNeurons
{
neuronsPrev.get(_targetSeq, _neurons[depth])!
.v[elem].delta +=
_coeffs[depth] * neurons.get(depth)!.v[elem].delta
neuronsPrev.get(_targetSeq, depth)!.v[elem].delta +=
neurons.get(depth)!.v[elem].delta
}
}

Expand Down Expand Up @@ -339,22 +259,18 @@ public class SelectNeuronsSeq: Layer1D

let pTargetSeq: [UInt32] = [UInt32(_targetSeq)]
let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
let pNbNeuronsPrev: [UInt32] = [UInt32(layerPrev.nbNeurons)]
let pNbBatch: [UInt32] = [UInt32(batchSize)]
let pSequence: [UInt32] = [UInt32(layerPrev.sequence)]

command = MetalKernel.get.createCommand(
"selectNeuronsSeqBackward", deviceID: deviceID
"selectSeqBackward", deviceID: deviceID
)
command.setBuffer(delta.metal, atIndex: 0)
command.setBytes(pNbNeurons, atIndex: 1)
command.setBytes(pNbNeuronsPrev, atIndex: 2)
command.setBytes(pTargetSeq, atIndex: 3)
command.setBuffer(_neuronsBuffer.metal, atIndex: 4)
command.setBuffer(_coeffsBuffer.metal, atIndex: 5)
command.setBytes(pNbBatch, atIndex: 6)
command.setBytes(pSequence, atIndex: 7)
command.setBuffer(layerPrev.delta.metal, atIndex: 8)
command.setBytes(pTargetSeq, atIndex: 2)
command.setBytes(pNbBatch, atIndex: 3)
command.setBytes(pSequence, atIndex: 4)
command.setBuffer(layerPrev.delta.metal, atIndex: 5)

command.dispatchThreads(
width: nbNeurons,
Expand Down
33 changes: 11 additions & 22 deletions Sources/GrAIdient/Metal/Kernel/LayerSeq.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,30 +98,25 @@ kernel void avgPoolSeqBackward(
}
}

kernel void selectNeuronsSeqForward(
kernel void selectSeqForward(
const device float * outsPrev,
constant uint * pNbNeurons,
constant uint * pNbNeuronsPrev,
constant uint * pTargetSeq,
constant uint * pNeurons,
constant float * pCoeffs,
constant uint * pNbBatch,
constant uint * pSequence,
device float * outs,
uint2 id [[ thread_position_in_grid ]])
{
uint targetSeq;
uint nbNeurons;
uint nbNeuronsPrev;
uint nbBatch;
uint sequence;

if (pTargetSeq && pNbNeurons && pNbNeuronsPrev && pNbBatch && pSequence &&
outsPrev && outs && pCoeffs && pNeurons)
if (pTargetSeq && pNbNeurons && pNbBatch && pSequence &&
outsPrev && outs)
{
targetSeq = *pTargetSeq;
nbNeurons = *pNbNeurons;
nbNeuronsPrev = *pNbNeuronsPrev;
nbBatch = *pNbBatch;
sequence = *pSequence;
}
Expand All @@ -137,36 +132,30 @@ kernel void selectNeuronsSeqForward(
}

uint offset = depth + nbNeurons * elem;
uint offsetPrev = pNeurons[depth] +
nbNeuronsPrev * targetSeq + sequence * nbNeuronsPrev * elem;
outs[offset] = pCoeffs[depth] * outsPrev[offsetPrev];
uint offsetPrev = depth +
nbNeurons * targetSeq + sequence * nbNeurons * elem;
outs[offset] = outsPrev[offsetPrev];
}

kernel void selectNeuronsSeqBackward(
kernel void selectSeqBackward(
const device float * delta,
constant uint * pNbNeurons,
constant uint * pNbNeuronsPrev,
constant uint * pTargetSeq,
constant uint * pNeurons,
constant float * pCoeffs,
constant uint * pNbBatch,
constant uint * pSequence,
device float * deltaPrev,
uint2 id [[ thread_position_in_grid ]])
{
uint nbNeurons;
uint nbNeuronsPrev;
uint nbBatch;
uint sequence;
uint targetSeq;

if (pNbNeurons && pNbNeuronsPrev &&
pTargetSeq && pNeurons && pCoeffs && pNbBatch && pSequence &&
if (pNbNeurons && pTargetSeq && pNbBatch && pSequence &&
deltaPrev && delta)
{
targetSeq = *pTargetSeq;
nbNeurons = *pNbNeurons;
nbNeuronsPrev = *pNbNeuronsPrev;
nbBatch = *pNbBatch;
sequence = *pSequence;
}
Expand All @@ -182,9 +171,9 @@ kernel void selectNeuronsSeqBackward(
}

uint offset = depth + nbNeurons * elem;
uint offsetPrev = pNeurons[depth] +
nbNeuronsPrev * targetSeq + sequence * nbNeuronsPrev * elem;
deltaPrev[offsetPrev] += pCoeffs[depth] * delta[offset];
uint offsetPrev = depth +
nbNeurons * targetSeq + sequence * nbNeurons * elem;
deltaPrev[offsetPrev] += delta[offset];
}

kernel void concat1SeqForward(
Expand Down
4 changes: 2 additions & 2 deletions Sources/GrAIdient/Metal/MetalKernel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -678,8 +678,8 @@ private class MetalDevice
"valueSeqForward",
"valueValueSeqBackward",
"valueScoreSeqBackward",
"selectNeuronsSeqForward",
"selectNeuronsSeqBackward",
"selectSeqForward",
"selectSeqBackward",
],
"Merge": [
"sum1",
Expand Down
2 changes: 1 addition & 1 deletion Sources/GrAIdient/Utils/Serialization.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ let LAYER_REGISTRY: [String: Codable.Type] = buildRegistry(
SumSeq.self,
SelectNeurons1D.self,
SelectNeurons2D.self,
SelectNeuronsSeq.self,
SelectSeq.self,
ValueSeq.self,
VQ2D.self,
])
Expand Down
Loading