Update

Author: junjihashimoto

README.md

@@ -1,18 +1,17 @@
 # Hasktorch Compose
 
-In hasktorch, model specifications, values, and inference are defined separately. there are cases combining commonly used models. For example, we might want to connect three linear layers, or add one linear layer to an existing model.
-This repo provides a library to easily compose existing models.
-We plan to provide both an untyped API and a typed API, but we will prioritize the development of the untyped API.
+In Hasktorch, model specifications, values, and inference are defined separately. This often necessitates combining commonly used models. For example, three linear layers may need to be connected in sequence, or a new linear layer could be added to an existing model. Hasktorch Compose provides a straightforward way to compose such models.
 
-This is an experimental library developed based on [hasktorch-skeleton](https://github.com/hasktorch/hasktorch-skeleton).
+In addition to simple model composition, this library aims to support extracting parts of models and sharing parameters between different models, such as ControlNet and RoLa. Both an untyped API and a typed API are planned, with initial development focused on the untyped API.
 
-List of planned features:
+Hasktorch Compose is an experimental library built on top of [hasktorch-skeleton](https://github.com/hasktorch/hasktorch-skeleton).
 
+**Planned Features:**
 - [x] Sequential
 - [ ] Extract layer
 - [ ] Test for each layer
 - [ ] Overlay layer
-- [ ] Concatenate layer
+- [x] Concatenate layer
 
 # Examples
 
@@ -30,11 +29,11 @@ mlpSpec =
   ReluSpec :>>:
   LinearSpec 64 32) :>>:
   ReluSpec :>>:
-  LinearSpec 32 10
+  LinearSpec 32 10 :>>:
+  LogSoftMaxSpec
 
 mlp :: (Randomizable MLPSpec MLP, HasForward MLP Tensor Tensor) => MLP -> Tensor -> Tensor
-mlp model input =
-  logSoftmax (Dim 1) $ forward model input
+mlp model input = forward model input
 ```
 
 ## Extract layer

src/Torch/Compose/NN.hs

@@ -117,6 +117,16 @@ instance Randomizable BatchNorm2dSpec BatchNorm2d where
     runningVar <- newMutableTensor $ ones' [channelSize]
     pure BatchNorm2d{..}
 
+data LogSoftMaxSpec = LogSoftMaxSpec deriving (Generic, Show, Eq)
+data LogSoftMax = LogSoftMax deriving (Generic, Parameterized, Show, Eq)
+instance Randomizable LogSoftMaxSpec LogSoftMax where
+  sample _ = pure LogSoftMax
+
+instance HasForward LogSoftMax Tensor Tensor where
+  forward _ = logSoftmax (Dim 1)
+  forwardStoch _ i = pure $ logSoftmax (Dim 1) i
+
+
 instance HasOutputs Linear Tensor where
   type Outputs Linear Tensor = Tensor
   toOutputs = forward
@@ -141,6 +151,19 @@ instance HasOutputShapes Relu Tensor where
   type OutputShapes Relu Tensor = [Int]
   toOutputShapes model a = shape $ forward model a
 
+instance HasOutputs LogSoftMax Tensor where
+  type Outputs LogSoftMax Tensor = Tensor
+  toOutputs = forward
+
+instance HasInputs LogSoftMax Tensor where
+  type Inputs LogSoftMax Tensor = Tensor
+  toInputs _ a = a
+
+instance HasOutputShapes LogSoftMax Tensor where
+  type OutputShapes LogSoftMax Tensor = [Int]
+  toOutputShapes model a = shape $ forward model a
+
+
 instance HasForwardAssoc Linear Tensor where
   type ForwardResult Linear Tensor = Tensor
   forwardAssoc = forward
@@ -149,3 +172,7 @@ instance HasForwardAssoc Relu Tensor where
   type ForwardResult Relu Tensor = Tensor
   forwardAssoc = forward
 
+instance HasForwardAssoc LogSoftMax Tensor where
+  type ForwardResult LogSoftMax Tensor = Tensor
+  forwardAssoc = forward
+

test/Spec.hs

@@ -24,16 +24,17 @@ import Torch.Compose.NN
 import Torch
 import GHC.Generics hiding ((:+:))
 
-newtype MLPSpec = MLPSpec (LinearSpec :>>: (ReluSpec :>>: (LinearSpec :>>: (ReluSpec :>>: LinearSpec)))) deriving (Generic)
-newtype MLP = MLP (Linear :>>: (Relu :>>: (Linear :>>: (Relu :>>: Linear)))) deriving (Generic)
+newtype MLPSpec = MLPSpec (LinearSpec :>>: ReluSpec :>>: LinearSpec :>>: ReluSpec :>>: LinearSpec :>>: LogSoftMaxSpec) deriving (Generic)
+newtype MLP = MLP (Linear :>>: Relu :>>: Linear :>>: Relu :>>: Linear :>>: LogSoftMax) deriving (Generic)
 
 mlpSpec :: MLPSpec
 mlpSpec = MLPSpec $
   LinearSpec 784 64 :>>:
   ReluSpec :>>:
   LinearSpec 64 32 :>>:
   ReluSpec :>>:
-  LinearSpec 32 10
+  LinearSpec 32 10 :>>:
+  LogSoftMaxSpec
 
 instance HasForward MLP Tensor Tensor where
   forward (MLP model) = forward model
@@ -43,8 +44,7 @@ instance Randomizable MLPSpec MLP where
   sample (MLPSpec spec) = MLP <$> sample spec
 
 mlp :: MLP -> Tensor -> Tensor
-mlp model input =
-  logSoftmax (Dim 1) $ forward model input
+mlp model input = forward model input
 
 main :: IO ()
 main = hspec $ do
@@ -58,10 +58,9 @@ main = hspec $ do
   it "Extract the last layer" $ do
     (MLP (model :: a)) <- sample mlpSpec
     let layer =  getLast model :: LastLayer a
-    shape layer.weight.toDependent `shouldBe` [10,32]
+    layer `shouldBe` LogSoftMax
   it "Extract all output shapes" $ do
     (MLP (model :: a)) <- sample mlpSpec
     let out = toOutputShapes model (ones' [2,784])
-        exp = [2,64] :>>: [2,64] :>>: [2,32] :>>: [2,32] :>>: [2,10]
+        exp = [2,64] :>>: [2,64] :>>: [2,32] :>>: [2,32] :>>: [2,10] :>>: [2,10]
     out `shouldBe` exp
-