hasktorch 0.2.1.3 → 0.2.1.4
raw patch · 5 files changed
+176/−1 lines, 5 files
Files
- hasktorch.cabal +2/−1
- src/Torch/Typed/Serialize.hs +58/−0
- src/Torch/Typed/Tensor.hs +21/−0
- test/Spec.hs +2/−0
- test/Torch/Typed/SerializeSpec.hs +93/−0
hasktorch.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: hasktorch-version: 0.2.1.3+version: 0.2.1.4 synopsis: Haskell bindings to libtorch, supporting both typed and untyped tensors. description: Hasktorch is a library for tensors and neural networks in Haskell. It is an independent open source community project which leverages the core C++ libraries shared by PyTorch. homepage: https://github.com/hasktorch/hasktorch#readme@@ -179,6 +179,7 @@ , Torch.Typed.NN.TransformerSpec , Torch.Typed.VisionSpec , Torch.Typed.NamedTensorSpec+ , Torch.Typed.SerializeSpec , SerializeSpec , RandomSpec , VisionSpec
src/Torch/Typed/Serialize.hs view
@@ -1,7 +1,9 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-} module Torch.Typed.Serialize where @@ -12,6 +14,9 @@ import qualified Torch.Internal.Type as ATen import qualified Torch.Tensor as D import Torch.Typed.Tensor+import Torch.Typed.Parameter+import Torch.Typed.NN+import Torch.Typed.Autograd -- | save list of tensors to file save ::@@ -32,3 +37,56 @@ FilePath -> IO (HList tensors) load = ATen.cast1 S.load++saveParameters ::+ forall model parameters tensors dtype device.+ ( Parameterized model,+ parameters ~ Parameters model,+ HMap' ToDependent parameters tensors,+ HMapM' IO MakeIndependent tensors parameters,+ HFoldrM IO TensorListFold [D.ATenTensor] tensors [D.ATenTensor],+ Apply TensorListUnfold [D.ATenTensor] (HUnfoldMRes IO [D.ATenTensor] tensors),+ HUnfoldM IO TensorListUnfold (HUnfoldMRes IO [D.ATenTensor] tensors) tensors+ ) =>+ model ->+ FilePath ->+ IO ()+saveParameters model filePath = save (hmap' ToDependent . flattenParameters $ model) filePath++loadParameters ::+ forall model parameters tensors dtype device.+ ( Parameterized model,+ parameters ~ Parameters model,+ HMap' ToDependent parameters tensors,+ HMapM' IO MakeIndependent tensors parameters,+ HFoldrM IO TensorListFold [D.ATenTensor] tensors [D.ATenTensor],+ Apply TensorListUnfold [D.ATenTensor] (HUnfoldMRes IO [D.ATenTensor] tensors),+ HUnfoldM IO TensorListUnfold (HUnfoldMRes IO [D.ATenTensor] tensors) tensors+ ) =>+ model ->+ FilePath ->+ IO model+loadParameters model filePath = do+ tensors <- load @tensors filePath+ params <- hmapM' MakeIndependent tensors+ pure $ replaceParameters model params++loadParametersWithSpec ::+ forall spec model parameters tensors dtype device.+ ( Randomizable spec model,+ Parameterized model,+ parameters ~ Parameters model,+ HMap' ToDependent parameters tensors,+ HMapM' IO MakeIndependent tensors parameters,+ HFoldrM IO TensorListFold [D.ATenTensor] tensors [D.ATenTensor],+ Apply TensorListUnfold [D.ATenTensor] (HUnfoldMRes IO [D.ATenTensor] tensors),+ HUnfoldM IO TensorListUnfold (HUnfoldMRes IO [D.ATenTensor] tensors) tensors+ ) =>+ spec ->+ FilePath ->+ IO model+loadParametersWithSpec spec filePath = do+ model <- sample spec+ tensors <- load @tensors filePath+ params <- hmapM' MakeIndependent tensors+ pure $ replaceParameters model params
src/Torch/Typed/Tensor.hs view
@@ -513,6 +513,27 @@ Tensor device dtype shape' selectIdx t idx = UnsafeMkTensor $ D.select (natValI @dim) (getFiniteI idx) (toDynamic t) +type family CheckIndexSelectDim (dim :: Nat) (shape :: [Nat]) (result :: Maybe [Nat]) :: [Nat] where+ CheckIndexSelectDim dim shape 'Nothing = TypeError (Text "Dim " :<>: ShowType dim :<>: Text " not found in shape " :<>: ShowType shape)+ CheckIndexSelectDim dim shape ('Just shape') = shape'++type IndexSelectDim (dim :: Nat) (shape :: [Nat]) (numIndices :: Nat) = CheckIndexSelectDim dim shape (ReplaceDim dim shape numIndices)++-- | Returns a new tensor which indexes the input tensor along dimension dim using the entries in index which is a tensor of datatype Int64.+-- The returned tensor has the same number of dimensions as the original tensor (input).+-- The dimth dimension has the same size as the length of index; other dimensions have the same size as in the original tensor.+-- +-- See https://pytorch.org/docs/stable/generated/torch.index_select.html for more information.+indexSelectDim ::+ forall (dim :: Nat) (shape :: [Nat]) (shape' :: [Nat]) (indexLength :: Nat) dtype device.+ ( KnownNat dim,+ shape' ~ IndexSelectDim dim shape indexLength+ ) =>+ Tensor device D.Int64 '[indexLength]+ -> Tensor device dtype shape+ -> Tensor device dtype shape'+indexSelectDim index inputs = UnsafeMkTensor $ D.indexSelect (natValI @dim) (toDynamic index) (toDynamic inputs)+ type family Numel (shape :: [Nat]) :: Nat where Numel '[] = 1 Numel (h ': t) = h * (Numel t)
test/Spec.hs view
@@ -37,6 +37,7 @@ import qualified Torch.Typed.TensorSpec0 import qualified Torch.Typed.TensorSpec1 import qualified Torch.Typed.VisionSpec+import qualified Torch.Typed.SerializeSpec main :: IO () main = hspec $ do@@ -76,4 +77,5 @@ Torch.Typed.TensorSpec0.spec Torch.Typed.TensorSpec1.spec Torch.Typed.VisionSpec.spec+ Torch.Typed.SerializeSpec.spec
+ test/Torch/Typed/SerializeSpec.hs view
@@ -0,0 +1,93 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}++module Torch.Typed.SerializeSpec+ ( Torch.Typed.SerializeSpec.spec,+ )+where++import Control.Monad (foldM)+import Data.Kind+import Data.Maybe+import Data.Proxy+import GHC.Exts (toList)+import GHC.Generics+import GHC.TypeLits+import Test.Hspec (Spec, describe, it, shouldBe)+import Test.QuickCheck ()+import Torch (ATenTensor)+import Torch.Internal.Class (Castable)+import Torch.Internal.Managed.Type.Context (manual_seed_L)+import Torch.Typed++import qualified Torch.DType as D+import qualified Torch.Device as D+++data+ MLPSpec+ (inputFeatures :: Nat)+ (outputFeatures :: Nat)+ (hiddenFeatures :: Nat)+ (dtype :: DType)+ (device :: (DeviceType, Nat))+ where+ MLPSpec ::+ forall inputFeatures outputFeatures hiddenFeatures dtype device.+ MLPSpec inputFeatures outputFeatures hiddenFeatures dtype device+ deriving (Show, Eq)++data+ MLP+ (inputFeatures :: Nat)+ (outputFeatures :: Nat)+ (hiddenFeatures :: Nat)+ (dtype :: D.DType)+ (device :: (D.DeviceType, Nat)) = MLP+ { layer0 :: Linear inputFeatures hiddenFeatures dtype device,+ layer1 :: Linear hiddenFeatures hiddenFeatures dtype device,+ layer2 :: Linear hiddenFeatures outputFeatures dtype device+ }+ deriving (Show, Generic, Parameterized)++instance+ ( KnownNat inputFeatures,+ KnownNat outputFeatures,+ KnownNat hiddenFeatures,+ KnownDType dtype,+ KnownDevice device,+ RandDTypeIsValid device dtype+ ) =>+ Randomizable+ (MLPSpec inputFeatures outputFeatures hiddenFeatures dtype device)+ (MLP inputFeatures outputFeatures hiddenFeatures dtype device)+ where+ sample _ =+ MLP+ <$> sample LinearSpec+ <*> sample LinearSpec+ <*> sample LinearSpec++saveMLP :: MLP 10 3 4 'D.Float '(D.CPU, 0) -> FilePath -> IO ()+saveMLP model filePath = saveParameters model filePath ++loadMLP :: MLP 10 3 4 'D.Float '(D.CPU, 0) -> FilePath -> IO (MLP 10 3 4 'D.Float '(D.CPU, 0))+loadMLP model filePath = loadParameters model filePath ++loadMLPWithSpec :: MLPSpec 10 3 4 'D.Float '(D.CPU, 0) -> FilePath -> IO (MLP 10 3 4 'D.Float '(D.CPU, 0))+loadMLPWithSpec spec filePath = loadParametersWithSpec spec filePath ++spec :: Spec+spec = pure ()+