diff --git a/examples/convnext/convnext.chpl b/examples/convnext/convnext.chpl new file mode 100644 index 000000000..ed1a7c6e3 --- /dev/null +++ b/examples/convnext/convnext.chpl @@ -0,0 +1,114 @@ +module ConvNext { + use Tensor; + use Network; + use List; + + class LayerNorm2d { + + } + + class ConvNextStem { + var conv: owned Conv2D(eltType); + var layernorm_2d: owned LayerNorm2d(eltType); + + proc init(in_channels: int, out_channels: int, kernel_size: int = 3) { + super.init(eltType); + this.conv = new Conv2D(eltType, channels=in_channels, features=out_channels, kernel=kernel_size, stride=kernel_size); + this.layernorm_2d = new LayerNorm2d(); // incomplete + } + + override proc forward(x: Tensor(eltType)): Tensor(eltType) { + x = this.conv(x) + x = this.layernorm_2d(x) // incomplete + return x + } + } + + class ConvNextBlock: Module(?) { + var depthwise_conv: owned Conv2D(eltType); // need groups param in convolutions + var layernorm_2d: owned LayerNorm2d(eltType); + var conv1: owned Conv2D(eltType); + // var gelu; // incomplete, need class wrappers for activation functions? + var conv2: owned Conv2D(eltType); + var layer_scale: owned Parameter(eltType); + + proc init(in_channels: int, layer_scale: int = 1e-6, kernel_size: int = 7, stochastic_depth_prob:int = 1) { + + /////////////////// need to finish this + super.init(eltType); + this.depthwise_conv = new Conv2D(); + this.conv1 = new Conv2D(); + // this.gelu + this.conv2 = new Conv2D(); + this.layer_scale = new Parameter(); + } + + override proc forward(x: Tensor(eltType)): Tensor(eltType) { + residual = x; + x = this.depthwise_conv(x); + x = this.layernorm_2d(x); + x = this.conv1(x); + x = x.gelu(); // should be changed to use classes + x = this.conv2(x); + x = this.layer_scale(x); + x = residual + x; + return x; + } + } + + class ConvNextDownsample { + var layernorm_2d: owned LayerNorm2d(eltType); + var conv: owned Conv2D(eltType); + + proc init(d_in: int, d_out: int, width:int = 2) { + super.init(eltType) + this.layer_norm = new LayerNorm2d(); // incomplete + this.conv = new Conv2D(eltType, channels=d_in, features=d_out, kernel=width, stride=width); + } + + override proc forward(x: Tensor(eltType)): Tensor(eltType) { + x = this.layernorm_2d(x); + x = this.conv(x); + return x; + } + } + + class ConvNextClassifier { + // var global_pool: // incomplete, need AdaptiveAvgPool2d (Global average pooling) + // adaptiveAvgPool2d is a method on NDArrays + // var flatten: // need Flatten + // var layer_norm: LayerNorm(); + // var linear: Linear(); + + proc init(d_in: int, d_out: int) { + // this.global_pool = AdaptiveAvgPool2d(...); + // this.flatten = Flatten(); + // this.layer_norm = LayerNorm(...); + // this.linear = Linear(); + } + + override proc forward(x: Tensor(eltType)): Tensor(eltType) { + x = this.global_pool(x); + x = this.flatten(x); + x = this.layer_norm(x); + x = this.linear(x); + return x; + } + } + + class ConvNext { + var layers: owned Sequential(eltType); + + proc init(in_channels: int, out_channels: int, blocks: [] int = [96]) { + var layer_list = list(owned Module(eltType)); + layer_list.pushBack()// etc + + + this.layers = new Sequential(layer_list); + } + + override proc forward(x: Tensor(eltType)): Tensor(eltType) { + x = layers(x); + } + } +} \ No newline at end of file