<commit_before>open Core_kernel.Std
open Tensorflow
module O = Ops

(* ResNet model for the mnist dataset.
   This is mostly a work in progess for now.

   Reference:
   - Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385)
*)
let label_count = Mnist_helper.label_count
let epochs = 10000
let batch_size = 1024

let hidden_nodes = 64

let scalar_tensor f =
  let array = Tensor.create1 Bigarray.float32 1 in
  Tensor.set array [| 0 |] f;
  array

let conv2d x ~out_features ~in_features ~stride:s ~kernel_size:k =
  let w = Var.normalf [ k; k; in_features; out_features ] ~stddev:0.1 in
  let b = Var.f [ out_features ] 0. in
  let conv = O.conv2D x w ~strides:[ 1; s; s; 1 ] ~padding:"SAME" in
  O.add conv b

let avg_pool x ~stride:s ~kernel_size:k =
  O.avgPool x ~ksize:[ 1; k; k; 1 ] ~strides:[ 1; s; s; 1 ] ~padding:"SAME"

let basic_block input_layer ~out_features ~in_features ~stride =
  let shortcut = conv2d input_layer ~out_features ~in_features ~stride ~kernel_size:1 in
  conv2d input_layer ~out_features ~in_features ~stride ~kernel_size:3
  |> O.relu
  |> conv2d ~out_features ~in_features:out_features ~stride:1 ~kernel_size:1
  |> O.add shortcut
  |> O.relu

let build_model ~xs =
  let layer = O.reshape xs (O.const_int ~type_:Int32 [ -1; 28; 28; 1 ]) in
  (* 3x3 convolution + max-pool. *)
  let layer = conv2d layer ~out_features:16 ~in_features:1 ~kernel_size:3 ~stride:1 in
  let layer = O.relu layer in

  let layer = basic_block layer ~out_features:16 ~in_features:16 ~stride:1 in

  let layer = avg_pool ~stride:1 ~kernel_size:8 layer in
  (* Final dense layer. *)
  let output_dim = 14*14*64 in
  let layer = O.reshape layer (O.const_int ~type_:Int32 [ -1; output_dim ]) in
  let w1 = Var.normalf [ output_dim; hidden_nodes ] ~stddev:0.1 in
  let b1 = Var.f [ hidden_nodes ] 0. in
  let w2 = Var.normalf [ hidden_nodes; label_count ] ~stddev:0.1 in
  let b2 = Var.f [ label_count ] 0. in
  O.(relu (layer *^ w1 + b1) *^ w2 + b2) |> O.softmax

let () =
  let keep_prob = O.placeholder [] ~type_:Float in
  let mnist = Mnist_helper.read_files () in
  let xs = O.placeholder [] ~type_:Float in
  let ys = O.placeholder [] ~type_:Float in
  let ys_node = O.Placeholder.to_node ys in
  let ys_ = build_model ~xs:(O.Placeholder.to_node xs) in
  let cross_entropy = O.(neg (reduce_mean (ys_node * log ys_))) in
  let accuracy =
    O.(equal (argMax ys_ one32) (argMax ys_node one32))<commit_msg>Add more substance to the resnet example (still missing batch normalization).<commit_after>open Core_kernel.Std
open Tensorflow
module O = Ops

(* ResNet model for the mnist dataset.
   This is mostly a work in progess for now.

   Reference:
   - Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385)
*)
let label_count = Mnist_helper.label_count
let epochs = 10000
let batch_size = 1024

let depth = 20
(* Number of hidden nodes in the final layer. *)
let hidden_nodes = 64

let scalar_tensor f =
  let array = Tensor.create1 Bigarray.float32 1 in
  Tensor.set array [| 0 |] f;
  array

let conv2d x ~out_features ~in_features ~stride:s ~kernel_size:k =
  let w = Var.normalf [ k; k; in_features; out_features ] ~stddev:0.1 in
  let b = Var.f [ out_features ] 0. in
  let conv = O.conv2D x w ~strides:[ 1; s; s; 1 ] ~padding:"SAME" in
  O.add conv b

let avg_pool x ~stride:s ~kernel_size:k =
  O.avgPool x ~ksize:[ 1; k; k; 1 ] ~strides:[ 1; s; s; 1 ] ~padding:"SAME"

let basic_block input_layer ~out_features ~in_features ~stride =
  let shortcut = conv2d input_layer ~out_features ~in_features ~stride ~kernel_size:1 in
  conv2d input_layer ~out_features ~in_features ~stride ~kernel_size:3
  |> O.relu
  |> conv2d ~out_features ~in_features:out_features ~stride:1 ~kernel_size:1
  |> O.add shortcut
  |> O.relu

let block_stack layer ~out_features ~in_features ~stride =
  let depth = (depth - 2) / 6 in
  let layer = basic_block layer ~out_features ~in_features ~stride in
  List.init (depth - 1) ~f:Fn.id
  |> List.fold ~init:layer ~f:(fun layer _idx ->
    basic_block layer ~out_features ~in_features:out_features ~stride:1)

let build_model ~xs =
  let layer = O.reshape xs (O.const_int ~type_:Int32 [ -1; 28; 28; 1 ]) in
  (* 3x3 convolution + max-pool. *)
  let layer = conv2d layer ~out_features:16 ~in_features:1 ~kernel_size:3 ~stride:1 in
  let layer = O.relu layer in

  let layer = block_stack layer ~out_features:16 ~in_features:16 ~stride:1 in
  let layer = block_stack layer ~out_features:32 ~in_features:16 ~stride:2 in
  let layer = block_stack layer ~out_features:64 ~in_features:32 ~stride:2 in

  let layer = avg_pool ~stride:1 ~kernel_size:8 layer in
  (* Final dense layer. *)
  let output_dim = 7*7*64 in
  let layer = O.reshape layer (O.const_int ~type_:Int32 [ -1; output_dim ]) in
  let w1 = Var.normalf [ output_dim; hidden_nodes ] ~stddev:0.1 in
  let b1 = Var.f [ hidden_nodes ] 0. in
  let w2 = Var.normalf [ hidden_nodes; label_count ] ~stddev:0.1 in
  let b2 = Var.f [ label_count ] 0. in
  O.(relu (layer *^ w1 + b1) *^ w2 + b2) |> O.softmax

let () =
  let keep_prob = O.placeholder [] ~type_:Float in
  let mnist = Mnist_helper.read_files () in
  let xs = O.placeholder [] ~type_:Float in
  let ys = O.placeholder [] ~type_:Float in
  let ys_node = O.Placeholder.to_node ys in
  let ys_ = build_model ~xs:(O.Placeholder.to_node xs) in
  let cross_entropy = O.(neg (reduce_mean (ys_node * log ys_))) in
  let accuracy =
    O.(equal (argMax ys_ one32) (argMax ys_node one32))
{'fraction_non_alphanumeric': 0.11538461538461539, 'fraction_numerical': 0.03645643456069996, 'mean_word_length': 3.842894969108561, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3492377', 'n_tokens_mistral': 2198, 'n_tokens_neox': 2038, 'n_words': 636}