gnn/gnn.go

package main

import (
	"fmt"
	"math"
	"math/rand"
	"sync"
)

const (
	Input   int     = 1433
	Hidden  int     = 64
	Output  int     = 7
	Sample  int     = 10000
	Batch   int     = 16
	Dropout float64 = 0.5
	Rate    float64 = 0.01
)

type (
	Parameter struct {
		W, B Matrix
	}

	Layer struct {
		d int
		f func(Matrix) Matrix
		p Parameter
		D Vector
	}
)

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 := A[i].Max()
		sum := 0.
		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 GetAggregation(G Graph, u, k int, l []Layer) Matrix {
	if len(G.A[u]) == 0 {
		return MakeMatrix(1, l[k].d)
	}
	// GCN
	A := MakeMatrix(1, l[k].d)
	for v := range G.A[u] {
		A.Add(Matrix{G.E[k][v]})
	}
	return A.Divide(float64(len(G.A[u])))
	// GAT
	// A := MakeMatrix(0, l[k].d)
	// for v := range G.A[u] {
	// 	A = append(A, G.E[k][v])
	// }
	// C := MakeMatrix(1, A.N())
	// Me := G.E[k][u].Modulus()
	// for i := 0; i < A.N(); i++ {
	// 	Ma := A[i].Modulus()
	// 	if Me > 0 && Ma > 0 {
	// 		C[0][i] = G.E[k][u].Dot(A[i]) / Me / Ma
	// 	}
	// }
	// return Softmax(C).Multiply(A)
}

func GetEmbedding(G Graph, u, k int, l []Layer, train bool) Matrix {
	E := MakeMatrix(1, l[k].d)
	if k == 0 {
		E.Add(Matrix{G.X[u]})
	} else {
		for v := range G.A[u] {
			GetEmbedding(G, v, k-1, l, train)
		}
		E.Add(GetAggregation(G, u, k-1, l).Multiply(l[k].p.W))
		E.Add(GetEmbedding(G, u, k-1, l, train).Multiply(l[k].p.B))
		l[k].f(E)
	}
	if train && l[k].D != nil {
		E.Dropout(l[k].D)
	}
	G.E[k][u] = E[0]
	return E
}

// A += B * C
func StartCalc(wg *sync.WaitGroup, A, B, C Matrix) {
	wg.Add(1)
	go func() {
		A.Add(B.Transpose().Multiply(C))
		wg.Done()
	}()
}

// A += B / c
func StartRefine(wg *sync.WaitGroup, A, B Matrix, c float64) {
	wg.Add(1)
	go func() {
		A.Add(B.Divide(c))
		wg.Done()
	}()
}

func Train(G Graph) []Layer {
	p1 := Parameter{MakeRandomMatrix(Input, Hidden), MakeRandomMatrix(Input, Hidden)}
	p2 := Parameter{MakeRandomMatrix(Hidden, Output), MakeRandomMatrix(Hidden, Output)}
	l := []Layer{{d: Input}, {d: Hidden, f: ReLU, p: p1}, {d: Output, f: Softmax, p: p2}}
	for i := 0; i < Sample; i++ {
		if i%100 == 0 {
			Test(G, l, false)
			// fmt.Println("sampling", i)
		}
		var wg sync.WaitGroup
		l[0].D, l[1].D = MakeDropoutVector(Input), MakeDropoutVector(Hidden)
		DW2, DB2 := MakeMatrix(Hidden, Output), MakeMatrix(Hidden, Output)
		DW1, DB1 := MakeMatrix(Input, Hidden), MakeMatrix(Input, Hidden)
		for j := 0; j < Batch; j++ {
			u := nodeId[rand.Intn(len(nodeId))]
			GetEmbedding(G, u, 2, l, true)
			delta := MakeMatrix(1, Output)
			delta[0][G.L[u]] = 1
			delta.Sub(Matrix{G.E[2][u]}).Divide(float64(Batch))
			StartCalc(&wg, DW2, GetAggregation(G, u, 1, l), delta)
			StartCalc(&wg, DB2, Matrix{G.E[1][u]}, delta)
			deltaB := delta.Multiply(l[2].p.B.Transpose())
			for k := 0; k < Hidden; k++ {
				if G.E[1][u][k] == 0 {
					deltaB[0][k] = 0
				}
			}
			StartCalc(&wg, DW1, GetAggregation(G, u, 0, l), deltaB)
			StartCalc(&wg, DB1, Matrix{G.E[0][u]}, deltaB)
			deltaW := delta.Multiply(l[2].p.W.Transpose())
			for v := range G.A[u] {
				delta = MakeMatrix(1, Hidden).Add(deltaW)
				for k := 0; k < Hidden; k++ {
					if G.E[1][v][k] == 0 {
						delta[0][k] = 0
					}
				}
				StartCalc(&wg, DW1, GetAggregation(G, v, 0, l), delta)
				StartCalc(&wg, DB1, Matrix{G.E[0][v]}, delta)
			}
			wg.Wait()
		}
		Rate := 0.2 * math.Exp(-float64(i)/1000)
		StartRefine(&wg, l[2].p.W, DW2, 1/Rate)
		StartRefine(&wg, l[2].p.B, DB2, 1/Rate)
		StartRefine(&wg, l[1].p.W, DW1, 1/Rate)
		StartRefine(&wg, l[1].p.B, DB1, 1/Rate)
		wg.Wait()
	}
	return l
}

func Test(G Graph, l []Layer, detail bool) {
	cnt1, cnt2, loss := 0, 0, 0.
	for u := range G.X {
		GetEmbedding(G, u, 2, l, false)
		id, _ := G.E[2][u].Max()
		if detail {
			fmt.Println(u, id)
		}
		if G.L[u] == id {
			cnt1++
		}
		if G.L[u] == id+Output {
			cnt2++
		}
		loss -= math.Log(G.E[2][u][G.L[u]%Output])
	}
	if detail {
		fmt.Println(cnt1, "/", len(nodeId), ",", cnt2, "/", Node-len(nodeId))
		fmt.Println(
			100*float64(cnt1)/float64(len(nodeId)), ",",
			100*float64(cnt2)/float64((Node-len(nodeId))), ",",
			loss/float64(len(G.X)),
		)
	} else {
		fmt.Println(100*float64(cnt2)/float64((Node-len(nodeId))), loss/float64(len(G.X)))
	}
}