@mineproness/tensorloop 2.2.3 → 2.2.5
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/package.json +1 -1
- package/src/Softmax/index.js +3 -3
package/package.json
CHANGED
package/src/Softmax/index.js
CHANGED
|
@@ -113,16 +113,16 @@ export class Softmax {
|
|
|
113
113
|
}
|
|
114
114
|
async crossEntropy(probs, target) {
|
|
115
115
|
const encoder = this.device.createCommandEncoder()
|
|
116
|
-
|
|
117
116
|
const tokenCount = probs.size / (this.vocabSize * 4)
|
|
118
117
|
const targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
|
|
119
118
|
if(!this.loss) this.loss = await this.gpu.createBuffer(4 * tokenCount)
|
|
120
119
|
if(!this.acc) this.acc = await this.gpu.createBuffer(4 * 1)
|
|
121
120
|
await this.gpu.WriteBuffer(targetBuffer, new Uint32Array(target))
|
|
122
|
-
const bindGroup1 = await this.gpu.createBindGroup(this.lossPipline, 0, [probs.mainBuffer, targetBuffer, this.loss])
|
|
123
|
-
const bindGroup2 = await this.gpu.createBindGroup(this.accPipline, 0 , [this.paramsForward, probs.mainBuffer,
|
|
121
|
+
const bindGroup1 = await this.gpu.createBindGroup(this.lossPipline, 0, [probs.mainBuffer, targetBuffer, this.loss, this.paramsForward])
|
|
122
|
+
const bindGroup2 = await this.gpu.createBindGroup(this.accPipline, 0 , [this.paramsForward, probs.mainBuffer, this.acc, targetBuffer])
|
|
124
123
|
await this.gpu.RunPipline(encoder, this.lossPipline, bindGroup1, Math.ceil((tokenCount / 4)))
|
|
125
124
|
await this.gpu.RunPipline(encoder, this.accPipline, bindGroup2, 1)
|
|
125
|
+
await this.gpu.submitQueue(encoder)
|
|
126
126
|
return {loss: await this.gpu.readBuffer(this.loss.size, this.loss), acc: (await this.gpu.readBuffer(this.acc.size, this.acc))[0]}
|
|
127
127
|
}
|
|
128
128
|
|