diff --git a/pkg/layers/flatten.go b/pkg/layers/flatten.go new file mode 100644 index 0000000..91f02d1 --- /dev/null +++ b/pkg/layers/flatten.go @@ -0,0 +1,77 @@ +package layers + +import ( + "fmt" + + "github.com/Hirogava/Go-NN-Learn/pkg/tensor" + "github.com/Hirogava/Go-NN-Learn/pkg/tensor/graph" +) + +type Flatten struct{} + +func NewFlatten() *Flatten { + return &Flatten{} +} + +func flattenOutputShape(shape []int) []int { + batch := shape[0] + features := 1 + for _, d := range shape[1:] { + features *= d + } + return []int{batch, features} +} + +func (f *Flatten) Forward(x *graph.Node) *graph.Node { + if x == nil || x.Value == nil { + panic("Flatten.Forward: input is nil") + } + if len(x.Value.Shape) < 2 { + panic(fmt.Sprintf("Flatten expects input with at least 2 dimensions, got %dD", len(x.Value.Shape))) + } + + inShape := append([]int{}, x.Value.Shape...) + outShape := flattenOutputShape(inShape) + + out, err := tensor.Reshape(x.Value, outShape) + if err != nil { + panic(fmt.Sprintf("Flatten.Forward: reshape failed: %v", err)) + } + + op := &flattenOp{ + x: x, + inShape: inShape, + outShape: outShape, + } + return graph.NewNode(out, []*graph.Node{x}, op) +} + +func (f *Flatten) Params() []*graph.Node { + return nil +} + +func (f *Flatten) Train() {} +func (f *Flatten) Eval() {} + +type flattenOp struct { + x *graph.Node + inShape []int + outShape []int +} + +func (op *flattenOp) Backward(grad *tensor.Tensor) { + gradIn, err := tensor.Reshape(grad, op.inShape) + if err != nil { + panic(fmt.Sprintf("Flatten.Backward: reshape failed: %v", err)) + } + + if op.x.Grad == nil { + op.x.Grad = gradIn + return + } + var addErr error + op.x.Grad, addErr = tensor.Add(op.x.Grad, gradIn) + if addErr != nil { + panic(fmt.Sprintf("Flatten.Backward: add grad failed: %v", addErr)) + } +} diff --git a/pkg/layers/flatten_test.go b/pkg/layers/flatten_test.go new file mode 100644 index 0000000..fb92b49 --- /dev/null +++ b/pkg/layers/flatten_test.go @@ -0,0 +1,180 @@ +package layers + +import ( + "reflect" + "testing" + + "github.com/Hirogava/Go-NN-Learn/pkg/autograd" + "github.com/Hirogava/Go-NN-Learn/pkg/tensor" + "github.com/Hirogava/Go-NN-Learn/pkg/tensor/graph" +) + +func TestFlattenForward_4D(t *testing.T) { + flatten := NewFlatten() + input := &graph.Node{ + Value: &tensor.Tensor{ + Data: []float64{ + 1, 2, + 3, 4, + 5, 6, + 7, 8, + }, + Shape: []int{1, 2, 2, 2}, + Strides: []int{8, 4, 2, 1}, + }, + } + + out := flatten.Forward(input) + + expectedShape := []int{1, 8} + if !reflect.DeepEqual(out.Value.Shape, expectedShape) { + t.Fatalf("expected shape %v, got %v", expectedShape, out.Value.Shape) + } + + expectedData := []float64{1, 2, 3, 4, 5, 6, 7, 8} + if !reflect.DeepEqual(out.Value.Data, expectedData) { + t.Fatalf("expected data %v, got %v", expectedData, out.Value.Data) + } +} + +func TestFlattenForward_3D(t *testing.T) { + flatten := NewFlatten() + input := &graph.Node{ + Value: &tensor.Tensor{ + Data: []float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, + Shape: []int{2, 3, 2}, + Strides: []int{6, 2, 1}, + }, + } + + out := flatten.Forward(input) + + expectedShape := []int{2, 6} + if !reflect.DeepEqual(out.Value.Shape, expectedShape) { + t.Fatalf("expected shape %v, got %v", expectedShape, out.Value.Shape) + } + + expectedData := []float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12} + if !reflect.DeepEqual(out.Value.Data, expectedData) { + t.Fatalf("expected data %v, got %v", expectedData, out.Value.Data) + } +} + +func TestFlattenForward_2DIdentity(t *testing.T) { + flatten := NewFlatten() + input := &graph.Node{ + Value: &tensor.Tensor{ + Data: []float64{1, 2, 3, 4}, + Shape: []int{2, 2}, + Strides: []int{2, 1}, + }, + } + + out := flatten.Forward(input) + + expectedShape := []int{2, 2} + if !reflect.DeepEqual(out.Value.Shape, expectedShape) { + t.Fatalf("expected shape %v, got %v", expectedShape, out.Value.Shape) + } + if !reflect.DeepEqual(out.Value.Data, input.Value.Data) { + t.Fatalf("expected data %v, got %v", input.Value.Data, out.Value.Data) + } +} + +func TestFlattenForward_Batch(t *testing.T) { + flatten := NewFlatten() + input := &graph.Node{ + Value: &tensor.Tensor{ + Data: []float64{ + 1, 2, 3, 4, + 5, 6, 7, 8, + }, + Shape: []int{2, 2, 2}, + Strides: []int{4, 2, 1}, + }, + } + + out := flatten.Forward(input) + + expectedShape := []int{2, 4} + if !reflect.DeepEqual(out.Value.Shape, expectedShape) { + t.Fatalf("expected shape %v, got %v", expectedShape, out.Value.Shape) + } +} + +func TestFlattenBackward(t *testing.T) { + autograd.SetGraph(autograd.NewGraph()) + defer autograd.ClearGraph() + + flatten := NewFlatten() + ctx := autograd.GetGraph() + input := ctx.RequireGrad(&tensor.Tensor{ + Data: []float64{ + 1, 2, + 3, 4, + 5, 6, + 7, 8, + }, + Shape: []int{1, 2, 2, 2}, + Strides: []int{8, 4, 2, 1}, + }) + + out := flatten.Forward(input) + if out.Operation == nil { + t.Fatal("Operation is nil") + } + + grad := &tensor.Tensor{ + Data: []float64{1, 1, 1, 1, 1, 1, 1, 1}, + Shape: []int{1, 8}, + Strides: []int{8, 1}, + } + out.Operation.Backward(grad) + + expectedGrad := []float64{1, 1, 1, 1, 1, 1, 1, 1} + if !reflect.DeepEqual(input.Grad.Data, expectedGrad) { + t.Fatalf("expected grad %v, got %v", expectedGrad, input.Grad.Data) + } +} + +func TestFlattenGradCheck4D(t *testing.T) { + proto := graph.NewNode(tensor.Zeros(1, 2, 2, 2), nil, nil) + + build := func(e *autograd.Engine, inputs []*graph.Node) *graph.Node { + return NewFlatten().Forward(inputs[0]) + } + + if !autograd.CheckGradientEngine(build, []*graph.Node{proto}, 1e-5, 1e-3) { + t.Error("grad check failed for Flatten 4D") + } +} + +func TestFlattenGradCheck3D(t *testing.T) { + proto := graph.NewNode(tensor.Zeros(2, 3, 2), nil, nil) + + build := func(e *autograd.Engine, inputs []*graph.Node) *graph.Node { + return NewFlatten().Forward(inputs[0]) + } + + if !autograd.CheckGradientEngine(build, []*graph.Node{proto}, 1e-5, 1e-3) { + t.Error("grad check failed for Flatten 3D") + } +} + +func TestFlattenInvalidInput(t *testing.T) { + flatten := NewFlatten() + input := &graph.Node{ + Value: &tensor.Tensor{ + Data: []float64{1, 2, 3}, + Shape: []int{3}, + Strides: []int{1}, + }, + } + + defer func() { + if r := recover(); r == nil { + t.Fatal("expected panic for 1D input") + } + }() + flatten.Forward(input) +} diff --git a/pkg/layers/gru.go b/pkg/layers/gru.go index d8b9d90..17fb570 100644 --- a/pkg/layers/gru.go +++ b/pkg/layers/gru.go @@ -250,8 +250,6 @@ func (op *gruOp) Backward(gradOutput *tensor.Tensor) { accumulate(op.x, dx) } -func sigmoid(v float64) float64 { return 1 / (1 + math.Exp(-v)) } - func splitGates(gates *tensor.Tensor, h int) (r, z, n *tensor.Tensor) { b := gates.Shape[0] r = gateSlice(gates, b, h, 0)