added more building block modules

CHANGELOG.md

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - fixed "nn/rnn NewLSTM() clashed weight names"
 - fixed some old API at `vision/aug/function.go`
 - fixed `tensor.OfSlice()` not supporting `[]int` data type
+- fixed make `tensor.ValueGo()` returning `[]int` instead of `[]int32`
+- added more building block modules: Dropout, MaxPool2D, Parameter, Identity
+- added nn.BatchNorm.Forward() with default training=true
 
 ## [Nofix]
 - ctype `long` caused compiling error in MacOS as noted on [#44]. Not working on linux box.

nn/batch-norm.go

@@ -90,3 +90,21 @@ func (bn *BatchNorm) ForwardT(xs *ts.Tensor, train bool) (retVal *ts.Tensor) {
 	return ts.MustBatchNorm(xs, bn.Ws, bn.Bs, bn.RunningMean, bn.RunningVar, train, bn.config.Momentum, bn.config.Eps, bn.config.CudnnEnable)
 
 }
+
+// Forward forwards inputs through the module.
+// NOTE.
+// This forwarding will update BatchNorm weight by default (training=true).
+// Wrap module with tensor.NoGrad() when running model inference mode.
+func (bn *BatchNorm) Forward(xs *ts.Tensor) (retVal *ts.Tensor) {
+	dim := xs.Dim()
+
+	if bn.Nd == 1 && dim != 2 && dim != 3 {
+		log.Fatalf("Expected an input tensor with 2 or 3 dims, got %v\n", xs.MustSize())
+	}
+
+	if bn.Nd > 1 && int(dim) != int(bn.Nd)+2 {
+		log.Fatalf("Expected an input tensor with %v dims, got %v\n", bn.Nd+2, xs.MustSize())
+	}
+
+	return ts.MustBatchNorm(xs, bn.Ws, bn.Bs, bn.RunningMean, bn.RunningVar, true, bn.config.Momentum, bn.config.Eps, bn.config.CudnnEnable)
+}

nn/other.go

@@ -0,0 +1,125 @@
+package nn
+
+import (
+	ts "github.com/sugarme/gotch/tensor"
+)
+
+// Dropout:
+// ========
+
+// Dropout represents a neural network dropout layer.
+type Dropout struct {
+	dropoutProb float64
+}
+
+// NewDropout creates a new Dropout layer
+func NewDropout(p float64) *Dropout {
+	return &Dropout{
+		dropoutProb: p,
+	}
+}
+
+// ForwardT implements ModuleT for Dropout layer.
+func (d *Dropout) ForwardT(input *ts.Tensor, train bool) (retVal *ts.Tensor) {
+	return ts.MustDropout(input, d.dropoutProb, train)
+}
+
+// NewParameter creates a kind of tensor that is considered as a module parameter.
+// Ref. https://pytorch.org/docs/stable/generated/torch.nn.parameter.Parameter.html
+func NewParameter(path *Path, name string, x *ts.Tensor, requireGradOpt ...bool) *ts.Tensor {
+	requiredGrad := true
+	if len(requireGradOpt) > 0 {
+		requiredGrad = requireGradOpt[0]
+	}
+
+	param := path.Add(name, x, requiredGrad)
+
+	return param
+}
+
+// Identity:
+// =========
+
+type Identity struct{}
+
+func (m *Identity) Forward(x *ts.Tensor) *ts.Tensor {
+	if x == nil {
+		return nil
+	}
+	return x.MustShallowClone()
+}
+
+func NewIdentity() *Identity {
+	return new(Identity)
+}
+
+// MaxPool2D:
+// ==========
+
+type MaxPool2D struct {
+	Kernel   []int64
+	Stride   []int64
+	Padding  []int64
+	Dilation []int64
+	CeilMode bool
+}
+
+type MaxPool2DOpts struct {
+	Stride   []int64
+	Padding  []int64
+	Dilation []int64
+	CeilMode bool
+}
+
+type MaxPool2DOpt func(*MaxPool2DOpts)
+
+func OptStrideMp2D(v []int64) MaxPool2DOpt {
+	return func(o *MaxPool2DOpts) {
+		o.Stride = v
+	}
+}
+
+func OptPaddingMp2D(v []int64) MaxPool2DOpt {
+	return func(o *MaxPool2DOpts) {
+		o.Padding = v
+	}
+}
+
+func OptDilationMp2D(v []int64) MaxPool2DOpt {
+	return func(o *MaxPool2DOpts) {
+		o.Dilation = v
+	}
+}
+
+func OptCeilModeMp2D(v bool) MaxPool2DOpt {
+	return func(o *MaxPool2DOpts) {
+		o.CeilMode = v
+	}
+}
+
+func DefaultMaxPool2DOpts() *MaxPool2DOpts {
+	return &MaxPool2DOpts{
+		Stride:   nil,
+		Padding:  []int64{0, 0},
+		Dilation: []int64{1, 1},
+	}
+}
+
+func NewMaxPool2D(kernelSize []int64, opts ...MaxPool2DOpt) *MaxPool2D {
+	o := DefaultMaxPool2DOpts()
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	return &MaxPool2D{
+		Kernel:   kernelSize,
+		Stride:   o.Stride,
+		Padding:  o.Padding,
+		Dilation: o.Dilation,
+		CeilMode: o.CeilMode,
+	}
+}
+
+func (m *MaxPool2D) Forward(x *ts.Tensor) *ts.Tensor {
+	return x.MustMaxPool2d(m.Kernel, m.Stride, m.Padding, m.Dilation, m.CeilMode, false)
+}