gnn/gnn.go

package main

import (
	"math"
)

const (
	HIDDEN int = 200
)

type (
	Parameter struct {
		Wo, Wi, B Matrix
	}

	Layer struct {
		d       int
		f       func(Matrix) Matrix
		p       Parameter
		O, I, E Matrix
	}
)

func ReLU(A Matrix) Matrix {
	for i := 0; i < A.N(); i++ {
		for j := 0; j < A.M(); j++ {
			A[i][j] = math.Max(0, A[i][j])
		}
	}
	return A
}

func Softmax(A Matrix) Matrix {
	for i := 0; i < A.N(); i++ {
		max, sum := 0., 0.
		for j := 0; j < A.M(); j++ {
			max = math.Max(max, A[i][j])
		}
		for j := 0; j < A.M(); j++ {
			A[i][j] = math.Exp(A[i][j] - max)
			sum += A[i][j]
		}
		for j := 0; j < A.M(); j++ {
			A[i][j] /= sum
		}
	}
	return A
}

func GetEmbedding(G Graph, u, k int, l []Layer) Vector {
	if k == 0 {
		return G.X[u]
	}
	l[k].E = Multiply(Matrix{GetEmbedding(G, u, k-1, l)}, l[k].p.B)
	l[k-1].O, l[k-1].I = MakeMatrix(1, l[k-1].d), MakeMatrix(1, l[k-1].d)
	Do, Di := 0, 0
	for v, w := range G.A[u] {
		if w == 1 {
			l[k-1].O.Add(Matrix{GetEmbedding(G, v, k-1, l)})
			Do++
		} else {
			l[k-1].I.Add(Matrix{GetEmbedding(G, v, k-1, l)})
			Di++
		}
	}
	if Do > 0 {
		l[k-1].O.Divide(float64(Do))
		l[k].E.Add(Multiply(l[k-1].O, l[k].p.Wo))
	}
	if Di > 0 {
		l[k-1].I.Divide(float64(Di))
		l[k].E.Add(Multiply(l[k-1].I, l[k].p.Wi))
	}
	return l[k].f(l[k].E)[0]
}

func Train(G Graph) []Layer {
	p1 := Parameter{MakeRandomMatrix(1433, HIDDEN), MakeRandomMatrix(1433, HIDDEN), MakeRandomMatrix(1433, HIDDEN)}
	p2 := Parameter{MakeRandomMatrix(HIDDEN, 7), MakeRandomMatrix(HIDDEN, 7), MakeRandomMatrix(HIDDEN, 7)}
	l := []Layer{{d: 1433}, {d: HIDDEN, f: ReLU, p: p1}, {d: 7, f: Softmax, p: p2}}
	for u, X := range G.X {
		GetEmbedding(G, u, 2, l)
		delta := MakeMatrix(1, 7)
		delta[0][nodeLabel[u]] = 1
		delta.Sub(l[2].E)
		DWo2, DWi2, DB2 := Multiply(l[1].O.Transpose(), delta), Multiply(l[1].I.Transpose(), delta), Multiply(l[1].E.Transpose(), delta)
		DWo2.Divide(10)
		DWi2.Divide(10)
		DB2.Divide(10)
		delta = Multiply(delta, l[2].p.B.Transpose())
		for i := 0; i < HIDDEN; i++ {
			if l[1].E[0][i] == 0 {
				delta[0][i] = 0
			}
		}
		DWo1, DWi1, DB1 := Multiply(l[0].O.Transpose(), delta), Multiply(l[0].I.Transpose(), delta), Multiply(Matrix{X}.Transpose(), delta)
		DWo1.Divide(10)
		DWi1.Divide(10)
		DB1.Divide(10)
		l[2].p.Wo.Add(DWo2)
		l[2].p.Wi.Add(DWi2)
		l[2].p.B.Add(DB2)
		l[1].p.Wo.Add(DWo1)
		l[1].p.Wi.Add(DWi1)
		l[1].p.B.Add(DB1)
	}
	return l
}