v0.17~preview.129.15+205

README.md

@@ -94,6 +94,7 @@ not work on macos).
 * [ResNet examples on CIFAR-10](./examples/cifar/README.md).
 * [Character-level RNN](./examples/char_rnn/README.md)
 * [Neural Style Transfer](./examples/neural_transfer/README.md)
+* [Reinforcement Learning](./examples/reinforcement-learning/README.md)
 
 Some more advanced applications from external repos:
 
@@ -107,7 +108,7 @@ Various pre-trained computer vision models are implemented in the vision library
 The weight files can be downloaded at the following links:
 
 
-* ResNet-18 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/resnet18.ot).
+* ResNet-18 [weights](https://github.com/LaurentMazare/tch-rs/releases/download/mw/resnet18.ot).
 * ResNet-34 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/resnet34.ot).
 * ResNet-50 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/resnet50.ot).
 * ResNet-101 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/resnet101.ot).
@@ -118,7 +119,7 @@ The weight files can be downloaded at the following links:
 * SqueezeNet 1.0 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/squeezenet1_0.ot).
 * SqueezeNet 1.1 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/squeezenet1_1.ot).
 * VGG-13 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/vgg13.ot).
-* VGG-16 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/vgg16.ot).
+* VGG-16 [weights](https://github.com/LaurentMazare/tch-rs/releases/download/mw/vgg16.ot).
 * AlexNet [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/alexnet.ot).
 * Inception-v3 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/inception-v3.ot).
 * MobileNet-v2 [weights](https://github.com/LaurentMazare/ocaml-torch/releases/download/v0.1-unstable/mobilenet-v2.ot).

bin/tensor_tools.ml

@@ -1,14 +1,14 @@
 open Base
 open Torch
-include Npy_lib
+open Npy
 
 let npz_tensors ~filename ~f =
-  let npz_file = Npy.Npz.open_in filename in
+  let npz_file = Npz.open_in filename in
   let named_tensors =
-    Npy.Npz.entries npz_file
-    |> List.map ~f:(fun tensor_name -> f tensor_name (Npy.Npz.read npz_file tensor_name))
+    Npz.entries npz_file
+    |> List.map ~f:(fun tensor_name -> f tensor_name (Npz.read npz_file tensor_name))
   in
-  Npy.Npz.close_in npz_file;
+  Npz.close_in npz_file;
   named_tensors
 ;;
 
@@ -20,7 +20,7 @@ let ls files =
       then
         npz_tensors ~filename ~f:(fun tensor_name packed_tensor ->
           match packed_tensor with
-          | Npy.P tensor ->
+          | P tensor ->
             let tensor_shape = Bigarray.Genarray.dims tensor |> Array.to_list in
             tensor_name, tensor_shape)
       else
@@ -36,7 +36,7 @@ let npz_to_pytorch npz_src pytorch_dst =
   let named_tensors =
     npz_tensors ~filename:npz_src ~f:(fun tensor_name packed_tensor ->
       match packed_tensor with
-      | Npy.P tensor ->
+      | P tensor ->
         (match Bigarray.Genarray.layout tensor with
          | Bigarray.C_layout -> tensor_name, Tensor.of_bigarray tensor
          | Bigarray.Fortran_layout -> failwith "fortran layout is not supported"))
@@ -62,19 +62,19 @@ let image_to_tensor image_src pytorch_dst resize =
 
 let pytorch_to_npz pytorch_src npz_dst =
   let named_tensors = Serialize.load_all ~filename:pytorch_src in
-  let npz_file = Npy.Npz.open_out npz_dst in
+  let npz_file = Npz.open_out npz_dst in
   List.iter named_tensors ~f:(fun (tensor_name, tensor) ->
     let write kind =
       let tensor = Tensor.to_bigarray tensor ~kind in
-      Npy.Npz.write npz_file tensor_name tensor
+      Npz.write npz_file tensor_name tensor
     in
     match Tensor.kind tensor with
     | T Float -> write Bigarray.float32
     | T Double -> write Bigarray.float64
     | T Int -> write Bigarray.int32
     | T Int64 -> write Bigarray.int64
     | _ -> Printf.failwithf "unsupported tensor kind for %s" tensor_name ());
-  Npy.Npz.close_out npz_file
+  Npz.close_out npz_file
 ;;
 
 let load_and_run_dummy_module device =

examples/basics/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names basics)
+ (libraries torch)
+ (preprocess
+  (pps ppx_jane)))

examples/char_rnn/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names char_rnn)
+ (libraries torch unix)
+ (preprocess
+  (pps ppx_jane)))

examples/cifar/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names cifar_train)
+ (libraries torch unix)
+ (preprocess
+  (pps ppx_jane)))

examples/gan/dune

@@ -0,0 +1,7 @@
+(executables
+ (modes byte exe)
+ (names began gan_stability mnist_cgan mnist_dcgan mnist_gan
+   progressive_growing_gan relativistic_dcgan)
+ (libraries torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

examples/jit/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names load_and_run)
+ (libraries torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

examples/min-gpt/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names mingpt)
+ (libraries torch unix)
+ (preprocess
+  (pps ppx_jane)))

examples/min-gpt/mingpt.ml

@@ -74,7 +74,7 @@ let block vs cfg =
     let ys =
       Layer.forward ln2 xs
       |> Layer.forward lin1
-      |> Tensor.gelu
+      |> Tensor.gelu ~approximate:"none"
       |> Layer.forward lin2
       |> Tensor.dropout ~p:resid_pdrop ~is_training
     in

examples/mnist/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names conv linear nn)
+ (libraries torch unix)
+ (preprocess
+  (pps ppx_jane)))

examples/neural_transfer/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names neural_transfer)
+ (libraries torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

examples/pretrained/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names finetuning predict)
+ (libraries torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

examples/reinforcement-learning/README.md

@@ -1,5 +1,12 @@
 # Reinforcement Learning Examples
 
+For the Python scripts, prepare your environment:
+
+```
+pip install gym
+pip install "gymnasium[atari, accept-rom-license]"
+```
+
 These examples illustrate how to implement a couple reinforcement learning
 algorithms to play Atari games.
 More details on the DQN examples can be found on this

examples/reinforcement-learning/a2c.ml

@@ -7,7 +7,6 @@
 *)
 open Base
 open Torch
-module E = Vec_env_gym_pyml
 
 let atari_game = "SpaceInvadersNoFrameskip-v4"
 let num_steps = 5
@@ -65,7 +64,7 @@ let train ~device =
     let dist_entropy =
       Tensor.(
         ~-(log_probs * probs)
-        |> sum_dim_intlist ~dim:[ -1 ] ~keepdim:false ~dtype:(T Float)
+        |> sum_dim_intlist ~dim:(Some [ -1 ]) ~keepdim:false ~dtype:(T Float)
         |> mean)
     in
     let advantages =
@@ -98,7 +97,7 @@ let valid ~filename ~device =
   let vs = Var_store.create ~frozen:true ~name:"a2c" () ~device in
   let model = model vs ~actions:action_space in
   Serialize.load_multi_ ~named_tensors:(Var_store.all_vars vs) ~filename;
-  let _ = Rollout.run rollout ~model in
+  let (_ : Rollout.rollout) = Rollout.run rollout ~model in
   let { Rollout.rewards = r; episodes = e } = Rollout.get_and_reset_totals rollout in
   Stdio.printf "%f (%.0f episodes)\n%!" (r /. e) e
 ;;

examples/reinforcement-learning/dqn_atari.ml

@@ -208,8 +208,8 @@ let preprocess () =
     let img =
       (* RGB to grey conversion. *)
       Tensor.(d 0 ~factor:0.299 + d 1 ~factor:0.587 + d 2 ~factor:0.114)
-      |> Tensor.slice ~dim:0 ~start:0 ~end_:210 ~step:2
-      |> Tensor.slice ~dim:1 ~start:0 ~end_:160 ~step:2
+      |> Tensor.slice ~dim:0 ~start:(Some 0) ~end_:(Some 210) ~step:2
+      |> Tensor.slice ~dim:1 ~start:(Some 0) ~end_:(Some 160) ~step:2
       |> Tensor.to_type ~type_:(T Uint8)
       |> Tensor.flip ~dims:[ 0; 1 ]
     in

examples/reinforcement-learning/dqn_pong.ml

@@ -176,8 +176,8 @@ let preprocess () =
       (* RGB to grey conversion. *)
       Tensor.(d 0 ~factor:0.299 + d 1 ~factor:0.587 + d 2 ~factor:0.114)
       |> Tensor.narrow ~dim:0 ~start:35 ~length:160
-      |> Tensor.slice ~dim:0 ~start:0 ~end_:160 ~step:2
-      |> Tensor.slice ~dim:1 ~start:0 ~end_:160 ~step:2
+      |> Tensor.slice ~dim:0 ~start:(Some 0) ~end_:(Some 160) ~step:2
+      |> Tensor.slice ~dim:1 ~start:(Some 0) ~end_:(Some 160) ~step:2
       |> Tensor.unsqueeze ~dim:0
     in
     let diff =

examples/reinforcement-learning/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names a2c dqn policy_gradient ppo dqn_atari dqn_pong)
+ (libraries pyml torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

examples/reinforcement-learning/policy_gradient.ml

@@ -103,7 +103,7 @@ let () =
       Tensor.(
         sum_dim_intlist
           (action_mask * log_softmax logits ~dim:1 ~dtype:(T Float))
-          ~dim:[ 1 ]
+          ~dim:(Some [ 1 ])
           ~keepdim:false
           ~dtype:(T Float))
     in

examples/reinforcement-learning/ppo.ml

@@ -18,7 +18,6 @@
 *)
 open Base
 open Torch
-module E = Vec_env_gym_pyml
 
 let atari_game = "SpaceInvadersNoFrameskip-v4"
 
@@ -94,7 +93,7 @@ let train ~device =
       let dist_entropy =
         Tensor.(
           ~-(log_probs * probs)
-          |> sum_dim_intlist ~dim:[ -1 ] ~keepdim:false ~dtype:(T Float)
+          |> sum_dim_intlist ~dim:(Some [ -1 ]) ~keepdim:false ~dtype:(T Float)
           |> mean)
       in
       let advantages = Tensor.( - ) (Tensor.to_device returns ~device) critic in
@@ -122,7 +121,7 @@ let valid ~filename ~device =
   let vs = Var_store.create ~frozen:true ~name:"a2c" () ~device in
   let model = model vs ~actions:action_space in
   Serialize.load_multi_ ~named_tensors:(Var_store.all_vars vs) ~filename;
-  let _ = Rollout.run rollout ~model in
+  let (_ : Rollout.rollout) = Rollout.run rollout ~model in
   let { Rollout.rewards = r; episodes = e } = Rollout.get_and_reset_totals rollout in
   Stdio.printf "%f (%.0f episodes)\n%!" (r /. e) e
 ;;

examples/translation/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names seq2seq)
+ (libraries torch unix)
+ (preprocess
+  (pps ppx_jane)))

examples/vae/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names vae)
+ (libraries torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

examples/yolo/darknet.ml

@@ -192,11 +192,9 @@ let detect xs ~image_height ~anchors ~classes ~device =
     |> Tensor.unsqueeze ~dim:0
   in
   slice_apply_and_set xs ~start:2 ~length:2 ~f:Tensor.(fun xs -> exp xs * anchors);
-  slice_apply_and_set
-    xs
-    ~start:0
-    ~length:4
-    ~f:Tensor.(fun xs -> xs * f (Float.of_int stride));
+  slice_apply_and_set xs ~start:0 ~length:4 ~f:(fun xs ->
+    let stride_float = Float.of_int stride in
+    Tensor.(xs * f stride_float));
   xs
 ;;
 

examples/yolo/dune

@@ -0,0 +1,6 @@
+(executables
+ (modes byte exe)
+ (names yolo)
+ (libraries torch torch_vision unix)
+ (preprocess
+  (pps ppx_jane)))

src/bindings/torch_bindings.ml

@@ -129,6 +129,8 @@ module C (F : Cstubs.FOREIGN) = struct
   end
 
   module Scalar = struct
+    let to_int64 = foreign "ats_to_int" (scalar @-> returning int64_t)
+    let to_float = foreign "ats_to_float" (scalar @-> returning float)
     let int = foreign "ats_int" (int64_t @-> returning scalar)
     let float = foreign "ats_float" (float @-> returning scalar)
     let free = foreign "ats_free" (scalar @-> returning void)

src/config/discover.ml

@@ -22,9 +22,20 @@ let or_else o ~f =
   | None -> f ()
 ;;
 
+let required_libs = [ "c10"; "torch_cpu"; "torch" ]
+
+let choose_dynamic_links ~lib_dir =
+  let optional_libs = [ "torch_cuda" ] in
+  required_libs
+  @ List.filter optional_libs ~f:(fun lib ->
+      match Sys_unix.file_exists [%string "%{lib_dir}/lib%{lib}.so"] with
+      | `Yes -> true
+      | _ -> false)
+;;
+
 let torch_flags () =
-  let dynamic_links = [ "c10"; "torch_cpu"; "torch" ] in
   let config ~include_dir ~lib_dir =
+    let dynamic_links = choose_dynamic_links ~lib_dir in
     let cflags =
       [ "-isystem"
       ; Printf.sprintf "%s" include_dir
@@ -34,9 +45,7 @@ let torch_flags () =
     in
     let libs =
       [ Printf.sprintf "-Wl,-rpath,%s" lib_dir; Printf.sprintf "-L%s" lib_dir ]
-      @ List.map
-          ~f:(fun lib -> [%string "-l%{lib_dir}"] /^ [%string "lib%{lib}.so"])
-          dynamic_links
+      @ List.map ~f:(fun lib -> [%string "-l%{lib}"]) dynamic_links
     in
     { C.Pkg_config.cflags; libs }
   in
@@ -67,6 +76,7 @@ let torch_flags () =
     |> or_else ~f:(fun () ->
          match Stdlib.Sys.getenv_opt "LIBTORCH_USE_SYSTEM" with
          | Some "1" ->
+           let dynamic_links = required_libs in
            Some
              { C.Pkg_config.cflags = []
              ; libs = List.map ~f:(fun lib -> [%string "-l%{lib}"]) dynamic_links

src/config/dune

@@ -1,6 +1,6 @@
 (executables
  (modes byte exe)
  (names discover)
- (libraries base dune-configurator)
+ (libraries base dune-configurator core_unix.sys_unix)
  (preprocess
   (pps ppx_jane)))

src/gen_bindings/gen.ml

@@ -704,7 +704,7 @@ let run ~declarations_filename ~gen_bindings ~gen_wrappers =
 
 let command =
   Command.basic
-    ~summary:"generate bindings or wrapper code for toch functions"
+    ~summary:"generate bindings or wrapper code for torch functions"
     (let%map_open.Command declarations_filename =
        flag "declarations" (required string) ~doc:"PATH path to Declarations.yaml"
      and gen_bindings =