diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog for hasktorch + +## 0.2.2.0 + +- Add support for Stackage lts-23.24 (GHC 9.8); the previous release used GHC 9.6. +- Add `BFloat16` to `DTypeIsFloatingPoint` and `KnownDType`. +- Implement the AdamW optimizer. +- Expose the parameter group interface for `CppOptim` (per-group Adam/AdamW options). +- Implement non-blocking host-to-device transfer. +- Remove partial signatures from `Torch.Typed.NN.Convolution`. diff --git a/hasktorch.cabal b/hasktorch.cabal --- a/hasktorch.cabal +++ b/hasktorch.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: hasktorch -version: 0.2.1.8 +version: 0.2.2.0 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 @@ -12,6 +12,7 @@ category: Codegen build-type: Custom extra-source-files: + CHANGELOG.md ./test/data/numpy_rawfile ./test/data/mnist-sample-labels-idx1-ubyte.gz ./test/data/mnist-sample-images-idx3-ubyte.gz @@ -107,7 +108,7 @@ cpp-options: -D__APPLE__ build-depends: async >= 2.2.5 && < 2.3 , base >= 4.7 && < 5 - , libtorch-ffi >= 2.0.1.10 && < 2.0.2 + , libtorch-ffi >= 2.0.2.0 && < 2.0.3 , libtorch-ffi-helper == 2.0.0.* , finite-typelits >= 0.1 && < 0.3 , ghc-typelits-extra >= 0.4.6 && < 0.6 diff --git a/src/Torch/Optim.hs b/src/Torch/Optim.hs --- a/src/Torch/Optim.hs +++ b/src/Torch/Optim.hs @@ -163,6 +163,74 @@ step = adam -- +-- AdamW +-- + +-- | State representation for AdamW Optimizer +data AdamW = AdamW + { beta1W :: Float, -- 1st moment forgetting factor + beta2W :: Float, -- 2nd moment forgetting factor + m1W :: [Tensor], -- 1st moment + m2W :: [Tensor], -- 2nd moment + iterW :: Int, -- iteration + weightDecayW :: Float -- weight decay + } + deriving (Show, Generic) + +instance NFData AdamW + +mkAdamW :: + Int -> + Float -> + Float -> + Float -> + [Parameter] -> + AdamW +mkAdamW iter beta1 beta2 weightDecay parameters = + AdamW + beta1 + beta2 + (initZeros <$> parameters) + (initZeros <$> parameters) + iter + weightDecay + where + initZeros = zerosLike . toDependent + +-- | AdamW step +adamw :: + -- | learning rate + LearningRate -> + -- | model parameter gradients + Gradients -> + -- | model parameters + [Tensor] -> + -- | adamw parameters + AdamW -> + -- | returns new parameters + updated adamw parameters + ([Tensor], AdamW) +adamw lr (Gradients gradients) parameters AdamW {..} = + (parameters', AdamW beta1W beta2W m1' m2' (iterW + 1) weightDecayW) + where + -- decaying averages of 1st & 2nd moments + f1 m1 dp = mulScalar beta1W m1 + mulScalar (1 - beta1W) dp + f2 m2 dp = mulScalar beta2W m2 + mulScalar (1 - beta2W) (dp * dp) + -- force to prevent spine laziness. See https://github.com/hasktorch/hasktorch/pull/728 + m1' = force $ zipWith f1 m1W gradients + m2' = force $ zipWith f2 m2W gradients + -- bias adjustment + a beta = divScalar (1 - beta ^ (iterW + 1)) + a1 = fmap (a beta1W) m1' + a2 = fmap (a beta2W) m2' + -- parameter update + eps = 1e-8 + update prevParam a1' a2' = prevParam - lr * (a1' / (sqrt a2' + eps) + mulScalar weightDecayW prevParam) + parameters' = zipWith3 update parameters a1 a2 + +instance Optimizer AdamW where + step = adamw + +-- -- Adagrad -- diff --git a/src/Torch/Optim/CppOptim.hs b/src/Torch/Optim/CppOptim.hs --- a/src/Torch/Optim/CppOptim.hs +++ b/src/Torch/Optim/CppOptim.hs @@ -241,3 +241,42 @@ loadState :: CppOptimizerState option -> FilePath -> IO () loadState (CppOptimizerState _ optimizer) file = cast2 LibTorch.load optimizer file + +data AdamwParamGroupOptions = AdamwParamGroupOptions + { adamwPgLr :: Double, + adamwPgBetas :: (Double, Double), + adamwPgEps :: Double, + adamwPgWeightDecay :: Double, + adamwPgAmsgrad :: Bool + } + deriving (Show, Eq) + +instance Default AdamwParamGroupOptions where + def = + AdamwParamGroupOptions + { adamwPgLr = 1e-3, + adamwPgBetas = (0.9, 0.999), + adamwPgEps = 1e-8, + adamwPgWeightDecay = 1e-2, + adamwPgAmsgrad = False + } + +initAdamwWithGroups :: AdamwParamGroupOptions -> [Tensor] -> [Tensor] -> IO (CppOptimizerState AdamwParamGroupOptions) +initAdamwWithGroups opt@AdamwParamGroupOptions {..} decayParams noDecayParams = do + v <- cast8 LibTorch.adamwWithParamGroups adamwPgLr (fst adamwPgBetas) (snd adamwPgBetas) adamwPgEps adamwPgWeightDecay adamwPgAmsgrad decayParams noDecayParams + return $ CppOptimizerState opt v + +cppOptimizerStepOnly :: CppOptimizerState option -> IO () +cppOptimizerStepOnly (CppOptimizerState _ optimizer) = cast1 LibTorch.stepOnly optimizer + +cppOptimizerZeroGrad :: CppOptimizerState option -> IO () +cppOptimizerZeroGrad (CppOptimizerState _ optimizer) = cast1 LibTorch.zeroGrad optimizer + +cppOptimizerSetGrads :: CppOptimizerState option -> [Tensor] -> IO () +cppOptimizerSetGrads (CppOptimizerState _ optimizer) grads = cast2 LibTorch.setParamGrads optimizer grads + +cppOptimizerGetAllParams :: CppOptimizerState option -> IO [Tensor] +cppOptimizerGetAllParams (CppOptimizerState _ optimizer) = cast1 LibTorch.getAllParams optimizer + +cppOptimizerSetLr :: CppOptimizerState option -> Double -> IO () +cppOptimizerSetLr (CppOptimizerState _ optimizer) lr = cast2 LibTorch.setLr optimizer lr diff --git a/src/Torch/Tensor.hs b/src/Torch/Tensor.hs --- a/src/Torch/Tensor.hs +++ b/src/Torch/Tensor.hs @@ -332,6 +332,58 @@ <> show di' <> "\"" +-- | Non-blocking variant of _toDevice for async H2D transfers. +-- IMPORTANT: Caller must synchronize before using the tensor (e.g., via +-- cudaStreamSynchronize or a synchronizing CUDA operation). +_toDeviceNonBlocking :: + -- | device to cast input to + Device -> + -- | input + Tensor -> + -- | output (may still be transferring) + Tensor +_toDeviceNonBlocking device' t = unsafePerformIO $ do + hasDevice <- case deviceType device' of + CPU -> pure True + CUDA -> cast0 ATen.hasCUDA + MPS -> cast0 ATen.hasMPS + let device = Torch.Tensor.device t + toDevice' + (deviceType device) + (deviceType device') + (deviceIndex device) + (deviceIndex device') + hasDevice + where + toDevice' dt dt' di di' _ | dt == dt' && di == di' = pure t + toDevice' CUDA CUDA di di' True | di /= di' = getOpts t >>= withDeviceIndex di' >>= toNonBlocking t + toDevice' CPU CUDA 0 di' True | di' >= 0 = getOpts t >>= withDeviceIndex di' >>= toNonBlocking t + toDevice' CUDA CPU di 0 True | di >= 0 = getOpts t >>= withDeviceType CPU >>= toNonBlocking t + toDevice' CPU MPS 0 0 True = getOpts t >>= withDeviceType MPS >>= toNonBlocking t + toDevice' MPS CPU 0 0 True = getOpts t >>= withDeviceType CPU >>= toNonBlocking t + toDevice' dt dt' di di' _ = + error $ + "cannot move tensor from \"" + <> show dt + <> ":" + <> show di + <> "\" to \"" + <> show dt' + <> ":" + <> show di' + <> "\"" + getOpts :: Tensor -> IO TensorOptions + getOpts = cast1 ATen.tensor_options + withDeviceType :: DeviceType -> TensorOptions -> IO TensorOptions + withDeviceType dt opts = cast2 ATen.tensorOptions_device_D opts dt + withDeviceIndex :: Int16 -> TensorOptions -> IO TensorOptions + withDeviceIndex di opts = cast2 ATen.tensorOptions_device_index_s opts di + toNonBlocking :: Tensor -> TensorOptions -> IO Tensor + toNonBlocking t opts = cast4 ATen.tensor_to_obb t opts nonBlocking copy + where + nonBlocking = True + copy = False + toDeviceWithTensor :: Tensor -> Tensor -> Tensor toDeviceWithTensor reference input = unsafePerformIO $ cast2 ATen.tensor_to_device reference input diff --git a/src/Torch/Typed/Auxiliary.hs b/src/Torch/Typed/Auxiliary.hs --- a/src/Torch/Typed/Auxiliary.hs +++ b/src/Torch/Typed/Auxiliary.hs @@ -313,6 +313,7 @@ DTypeIsFloatingPoint _ 'D.Half = () DTypeIsFloatingPoint _ 'D.Float = () DTypeIsFloatingPoint _ 'D.Double = () + DTypeIsFloatingPoint _ 'D.BFloat16 = () DTypeIsFloatingPoint '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype type family DTypeIsIntegral (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where diff --git a/src/Torch/Typed/NN/Convolution.hs b/src/Torch/Typed/NN/Convolution.hs --- a/src/Torch/Typed/NN/Convolution.hs +++ b/src/Torch/Typed/NN/Convolution.hs @@ -8,7 +8,6 @@ {-# LANGUAGE FunctionalDependencies #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE KindSignatures #-} -{-# LANGUAGE PartialTypeSignatures #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -16,7 +15,6 @@ {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE OverloadedRecordDot #-} {-# LANGUAGE DuplicateRecordFields #-} -{-# OPTIONS_GHC -Wno-partial-type-signatures #-} module Torch.Typed.NN.Convolution where @@ -64,14 +62,24 @@ deriving (Show, Generic, Parameterized) -- | conv1d --- The constraints on this one are _very_ involved, so the partial signatures --- make the code significantly cleaner. conv1dForward :: - forall stride padding. - _ => - Conv1d _ _ _ _ _ -> - Tensor _ _ _ -> - Tensor _ _ _ + forall stride padding inputChannelSize outputChannelSize kernelSize inputSize batchSize outputSize dtype device. + ( All + KnownNat + '[ stride, + padding, + inputChannelSize, + outputChannelSize, + kernelSize, + inputSize, + batchSize, + outputSize + ], + ConvSideCheck inputSize kernelSize stride padding outputSize + ) => + Conv1d inputChannelSize outputChannelSize kernelSize dtype device -> + Tensor device dtype '[batchSize, inputChannelSize, inputSize] -> + Tensor device dtype '[batchSize, outputChannelSize, outputSize] conv1dForward Conv1d {..} input = conv1d @stride @padding (toDependent weight) @@ -147,14 +155,30 @@ deriving (Show, Generic, Parameterized) -- | conv2d --- The constraints on this one are _very_ involved, so the partial signatures --- make the code significantly cleaner. conv2dForward :: - forall stride padding. - _ => - Conv2d _ _ _ _ _ _ -> - Tensor _ _ _ -> - Tensor _ _ _ + forall stride padding inputChannelSize outputChannelSize kernelSize0 kernelSize1 inputSize0 inputSize1 batchSize outputSize0 outputSize1 dtype device. + ( All + KnownNat + '[ Torch.Typed.Auxiliary.Fst stride, + Torch.Typed.Auxiliary.Snd stride, + Torch.Typed.Auxiliary.Fst padding, + Torch.Typed.Auxiliary.Snd padding, + inputChannelSize, + outputChannelSize, + kernelSize0, + kernelSize1, + inputSize0, + inputSize1, + batchSize, + outputSize0, + outputSize1 + ], + ConvSideCheck inputSize0 kernelSize0 (Torch.Typed.Auxiliary.Fst stride) (Torch.Typed.Auxiliary.Fst padding) outputSize0, + ConvSideCheck inputSize1 kernelSize1 (Torch.Typed.Auxiliary.Snd stride) (Torch.Typed.Auxiliary.Snd padding) outputSize1 + ) => + Conv2d inputChannelSize outputChannelSize kernelSize0 kernelSize1 dtype device -> + Tensor device dtype '[batchSize, inputChannelSize, inputSize0, inputSize1] -> + Tensor device dtype '[batchSize, outputChannelSize, outputSize0, outputSize1] conv2dForward Conv2d {..} input = conv2d @stride @padding (toDependent weight) @@ -240,14 +264,33 @@ deriving (Show, Generic, Parameterized) -- | conv3d --- The constraints on this one are _very_ involved, so the partial signatures --- make the code significantly cleaner. conv3dForward :: - forall stride padding. - _ => - Conv3d _ _ _ _ _ _ _ -> - Tensor _ _ _ -> - Tensor _ _ _ + forall stride padding inputChannelSize outputChannelSize kernelSize0 kernelSize1 kernelSize2 inputSize0 inputSize1 inputSize2 batchSize outputSize0 outputSize1 outputSize2 dtype device. + ( All + KnownNat + '[ Fst3 stride, + Snd3 stride, + Trd3 stride, + Fst3 padding, + Snd3 padding, + Trd3 padding, + inputChannelSize, + outputChannelSize, + kernelSize0, + kernelSize1, + kernelSize2, + inputSize0, + inputSize1, + inputSize2, + batchSize + ], + ConvSideCheck inputSize0 kernelSize0 (Fst3 stride) (Fst3 padding) outputSize0, + ConvSideCheck inputSize1 kernelSize1 (Snd3 stride) (Snd3 padding) outputSize1, + ConvSideCheck inputSize2 kernelSize2 (Trd3 stride) (Trd3 padding) outputSize2 + ) => + Conv3d inputChannelSize outputChannelSize kernelSize0 kernelSize1 kernelSize2 dtype device -> + Tensor device dtype '[batchSize, inputChannelSize, inputSize0, inputSize1, inputSize2] -> + Tensor device dtype '[batchSize, outputChannelSize, outputSize0, outputSize1, outputSize2] conv3dForward Conv3d {..} input = conv3d @stride @padding (toDependent weight) @@ -331,14 +374,24 @@ deriving (Show, Generic, Parameterized) -- | convTranspose1d --- The constraints on this one are _very_ involved, so the partial signatures --- make the code significantly cleaner. convTranspose1dForward :: - forall stride padding. - _ => - ConvTranspose1d _ _ _ _ _ -> - Tensor _ _ _ -> - Tensor _ _ _ + forall stride padding inputChannelSize outputChannelSize kernelSize inputSize batchSize outputSize dtype device. + ( All + KnownNat + '[ stride, + padding, + inputChannelSize, + outputChannelSize, + kernelSize, + inputSize, + batchSize, + outputSize + ], + ConvSideCheck inputSize kernelSize stride padding outputSize + ) => + ConvTranspose1d inputChannelSize outputChannelSize kernelSize dtype device -> + Tensor device dtype '[batchSize, inputChannelSize, inputSize] -> + Tensor device dtype '[batchSize, outputChannelSize, outputSize] convTranspose1dForward ConvTranspose1d {..} input = convTranspose1d @stride @padding (toDependent weight) @@ -414,14 +467,30 @@ deriving (Show, Generic, Parameterized) -- | convTranspose2d --- The constraints on this one are _very_ involved, so the partial signatures --- make the code significantly cleaner. convTranspose2dForward :: - forall stride padding. - _ => - ConvTranspose2d _ _ _ _ _ _ -> - Tensor _ _ _ -> - Tensor _ _ _ + forall stride padding inputChannelSize outputChannelSize kernelSize0 kernelSize1 inputSize0 inputSize1 batchSize outputSize0 outputSize1 dtype device. + ( All + KnownNat + '[ Torch.Typed.Auxiliary.Fst stride, + Torch.Typed.Auxiliary.Snd stride, + Torch.Typed.Auxiliary.Fst padding, + Torch.Typed.Auxiliary.Snd padding, + inputChannelSize, + outputChannelSize, + kernelSize0, + kernelSize1, + inputSize0, + inputSize1, + batchSize, + outputSize0, + outputSize1 + ], + ConvSideCheck inputSize0 kernelSize0 (Torch.Typed.Auxiliary.Fst stride) (Torch.Typed.Auxiliary.Fst padding) outputSize0, + ConvSideCheck inputSize1 kernelSize1 (Torch.Typed.Auxiliary.Snd stride) (Torch.Typed.Auxiliary.Snd padding) outputSize1 + ) => + ConvTranspose2d inputChannelSize outputChannelSize kernelSize0 kernelSize1 dtype device -> + Tensor device dtype '[batchSize, inputChannelSize, inputSize0, inputSize1] -> + Tensor device dtype '[batchSize, outputChannelSize, outputSize0, outputSize1] convTranspose2dForward ConvTranspose2d {..} input = convTranspose2d @stride @padding (toDependent weight) @@ -507,14 +576,33 @@ deriving (Show, Generic, Parameterized) -- | convTranspose3d --- The constraints on this one are _very_ involved, so the partial signatures --- make the code significantly cleaner. convTranspose3dForward :: - forall stride padding. - _ => - ConvTranspose3d _ _ _ _ _ _ _ -> - Tensor _ _ _ -> - Tensor _ _ _ + forall stride padding inputChannelSize outputChannelSize kernelSize0 kernelSize1 kernelSize2 inputSize0 inputSize1 inputSize2 batchSize outputSize0 outputSize1 outputSize2 dtype device. + ( All + KnownNat + '[ Fst3 stride, + Snd3 stride, + Trd3 stride, + Fst3 padding, + Snd3 padding, + Trd3 padding, + inputChannelSize, + outputChannelSize, + kernelSize0, + kernelSize1, + kernelSize2, + inputSize0, + inputSize1, + inputSize2, + batchSize + ], + ConvSideCheck inputSize0 kernelSize0 (Fst3 stride) (Fst3 padding) outputSize0, + ConvSideCheck inputSize1 kernelSize1 (Snd3 stride) (Snd3 padding) outputSize1, + ConvSideCheck inputSize2 kernelSize2 (Trd3 stride) (Trd3 padding) outputSize2 + ) => + ConvTranspose3d inputChannelSize outputChannelSize kernelSize0 kernelSize1 kernelSize2 dtype device -> + Tensor device dtype '[batchSize, inputChannelSize, inputSize0, inputSize1, inputSize2] -> + Tensor device dtype '[batchSize, outputChannelSize, outputSize0, outputSize1, outputSize2] convTranspose3dForward ConvTranspose3d {..} input = convTranspose3d @stride @padding (toDependent weight) diff --git a/src/Torch/Typed/Tensor.hs b/src/Torch/Typed/Tensor.hs --- a/src/Torch/Typed/Tensor.hs +++ b/src/Torch/Typed/Tensor.hs @@ -93,6 +93,9 @@ instance KnownDType 'D.Half where dtypeVal = D.Half +instance KnownDType 'D.BFloat16 where + dtypeVal = D.BFloat16 + instance KnownDType 'D.Float where dtypeVal = D.Float