From 4c57bd183d54428484603d310d28a611a05bf893 Mon Sep 17 00:00:00 2001
From: itsubaki <1759459+itsubaki@users.noreply.github.com>
Date: Sun, 22 Oct 2023 01:14:19 +0900
Subject: [PATCH] Update some files

---
 cmd/lstm/main.go | 65 ++++++++++++++++++++++++++++--------------------
 1 file changed, 38 insertions(+), 27 deletions(-)

diff --git a/cmd/lstm/main.go b/cmd/lstm/main.go
index e532d89..4e4ba95 100644
--- a/cmd/lstm/main.go
+++ b/cmd/lstm/main.go
@@ -13,15 +13,37 @@ import (
 	"github.com/itsubaki/autograd/vector"
 )
 
-type SinCurve struct {
-	N         int
+type DataLoader struct {
 	BatchSize int
+	N         int
 	Data      []float64
 	Label     []float64
 	iter      int
 }
 
-func NewSinCurve(batchSize int) *SinCurve {
+func (d *DataLoader) Next() bool {
+	next := (d.iter+1)*d.BatchSize < d.N
+	if !next {
+		d.iter = 0
+	}
+
+	return next
+}
+
+func (d *DataLoader) Batch() (*variable.Variable, *variable.Variable) {
+	begin, end := d.iter*d.BatchSize, (d.iter+1)*d.BatchSize
+	x, y := vector.Transpose(d.Data[begin:end]), vector.Transpose(d.Label[begin:end])
+	d.iter++
+	return variable.NewOf(x...), variable.NewOf(y...)
+}
+
+type SinCurve struct {
+	N     int
+	Data  []float64
+	Label []float64
+}
+
+func NewSinCurve() *SinCurve {
 	N, noise := 1000, 0.05
 
 	x := make([]float64, N)
@@ -35,28 +57,10 @@ func NewSinCurve(batchSize int) *SinCurve {
 	}
 
 	return &SinCurve{
-		N:         N,
-		BatchSize: batchSize,
-		Data:      y[:len(x)-1],
-		Label:     y[1:],
-		iter:      0,
-	}
-}
-
-func (d *SinCurve) Next() bool {
-	next := (d.iter+1)*d.BatchSize < d.N
-	if !next {
-		d.iter = 0
+		N:     N,
+		Data:  y[:len(x)-1],
+		Label: y[1:],
 	}
-
-	return next
-}
-
-func (d *SinCurve) Read() (*variable.Variable, *variable.Variable) {
-	begin, end := d.iter*d.BatchSize, (d.iter+1)*d.BatchSize
-	x, y := vector.Transpose(d.Data[begin:end]), vector.Transpose(d.Label[begin:end])
-	d.iter++
-	return variable.NewOf(x...), variable.NewOf(y...)
 }
 
 func main() {
@@ -67,16 +71,23 @@ func main() {
 	flag.IntVar(&bpttLength, "bptt-length", 30, "")
 	flag.Parse()
 
+	dataset := NewSinCurve()
+	dataloader := &DataLoader{
+		BatchSize: batchSize,
+		N:         dataset.N,
+		Data:      dataset.Data,
+		Label:     dataset.Label,
+	}
+
 	m := model.NewLSTM(hiddenSize, 1)
 	o := optimizer.SGD{LearningRate: 0.01}
 
-	dataset := NewSinCurve(batchSize)
 	for i := 0; i < epoch; i++ {
 		m.ResetState()
 
 		loss, count := variable.Const(0), 0
-		for dataset.Next() {
-			x, t := dataset.Read()
+		for dataloader.Next() {
+			x, t := dataloader.Batch()
 			y := m.Forward(x)
 			loss = F.Add(loss, F.MeanSquaredError(y, t))