golang 实现rnn

package main

import (
    "fmt"
    "math"
)

func sigmoid(x float64) float64 {
    //sigmoid 激活函数
    return 1 / (1 + math.Exp(-x))
}

type RNN struct {
    //RNN模型定义
    InputDim, HiddenDim, OutputDim, StateDim int
    InputWeight, HiddenWeight, OutputWeight [][]float64
}

func NewRNN(inputDim, hiddenDim, outputDim, stateDim int) *RNN {
    rnn := &RNN{}
    rnn.InputDim = inputDim
    rnn.HiddenDim = hiddenDim
    rnn.OutputDim = outputDim
    rnn.StateDim = stateDim
    rnn.InputWeight = make([][]float64, inputDim)
    for i := range rnn.InputWeight {
        rnn.InputWeight[i] = make([]float64, hiddenDim)
    }
    rnn.HiddenWeight = make([][]float64, hiddenDim)
    for i := range rnn.HiddenWeight {
        rnn.HiddenWeight[i] = make([]float64, hiddenDim)
    }
    rnn.OutputWeight = make([][]float64, hiddenDim)
    for i := range rnn.OutputWeight {
        rnn.OutputWeight[i] = make([]float64, outputDim)
    }
    return rnn
}

func (rnn *RNN) Forward(input []float64) ([]float64, [][]float64) {
    h := make([]float64, rnn.HiddenDim)
    state := make([]float64, rnn.StateDim)
    output := make([]float64, rnn.OutputDim)
    //前向传播
    for i := 0; i < rnn.HiddenDim; i++ {
        for j := 0; j < rnn.InputDim; j++ {
            h[i] += input[j] * rnn.InputWeight[j][i]
        }
        for j := 0; j < rnn.HiddenDim; j++ {
            h[i] += state[j] * rnn.HiddenWeight[j][i]
        }
        h[i] = sigmoid(h[i])
    }
    for i := 0; i < rnn.OutputDim; i++ {
        for j := 0; j < rnn.HiddenDim; j++ {
            output[i] += h[j] * rnn.OutputWeight[j][i]
        }
    }
    return output, [][]float64{nil, nil, nil}
}

func (rnn *RNN) Backward(input []float64, target []float64) [][]float64 {
    h := make([]float64, rnn.HiddenDim)
    state := make([]float64, rnn.StateDim)
    output := make([]float64, rnn.OutputDim)
    delta := make([]float64, rnn.OutputDim)
    deltaH := make([]float64, rnn.HiddenDim)
    //计算损失
    loss := 0.0
    for i := 0; i < rnn.OutputDim; i++ {
        loss += math.Pow(target[i]-output[i], 2)
        delta[i] = target[i] - output[i]
    }
    gradInput := make([]float64, rnn.InputDim)
    gradInputWeight := make([][]float64, rnn.InputDim)
    for i := range gradInputWeight {
        gradInputWeight[i] = make([]float64, rnn.HiddenDim)
    }
    gradHiddenWeight := make([][]float64, rnn.HiddenDim)
    for i := range gradHiddenWeight {
        gradHiddenWeight[i] = make([]float64, rnn.HiddenDim)
    }
    gradOutputWeight := make([][]float64, rnn.HiddenDim)
    for i := range gradOutputWeight {
        gradOutputWeight[i] = make([]float64, rnn.OutputDim)
    }
    //反向传播
    for i := 0; i < rnn.OutputDim; i++ {
        for j := 0; j < rnn.HiddenDim; j++ {
            gradOutputWeight[j][i] = h[j] * delta[i]
            deltaH[j] += delta[i] * rnn.OutputWeight[j][i]
        }
    }
    for i := 0; i < rnn.HiddenDim; i++ {
        deltaH[i] *= h[i] * (1 - h[i])
        for j := 0; j < rnn.HiddenDim; j++ {
            gradHiddenWeight[j][i] = state[j] * deltaH[i]
            if i == 0 {
                gradInput[j] = input[j] * deltaH[0]
                for k := 0; k < rnn.HiddenDim; k++ {
                    gradInputWeight[j][k] = input[j] * deltaH[0] * h[k]
                }
            }
        }
        for j := 0; j < rnn.StateDim; j++ {
            state[j] = deltaH[i] * rnn.HiddenWeight[j][i]
        }
    }
    return [][]float64{gradInput, gradInputWeight, gradHiddenWeight, gradOutputWeight}
}

func main() {
    //定义RNN模型
    rnn := NewRNN(2, 2, 1, 2)
    rnn.InputWeight[0][0] = 0.5
    rnn.InputWeight[0][1] = 0.2
    rnn.InputWeight[1][0] = 0.1
    rnn.InputWeight[1][1] = 0.3
    rnn.HiddenWeight[0][0] = 0.4
    rnn.HiddenWeight[0][1] = 0.4
    rnn.HiddenWeight[1][0] = 0.5
    rnn.HiddenWeight[1][1] = 0.5
    rnn.OutputWeight[0][0] = 0.6
    rnn.OutputWeight[1][0] = 0.7
    //前向传播和反向传播
    output, _ := rnn.Forward([]float64{0.2, 0.4})
    fmt.Println("Output:", output)
    grad := rnn.Backward([]float64{0.2, 0.4}, []float64{0.9})
    fmt.Println("Grad:", grad)
}