1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
open Base
open Torch
let sub = Var_store.sub
let subi = Var_store.subi
let relu6 xs = Tensor.(min (relu xs) (f 6.))
let cbr vs ksize ~stride ~input_dim ?(groups = 1) output_dim =
let ksize, padding =
match ksize with
| `k1 -> 1, 0
| `k3 -> 3, 1
in
let conv =
Layer.conv2d_
(subi vs 0)
~ksize
~stride
~padding
~groups
~use_bias:false
~input_dim
output_dim
in
let bn = Layer.batch_norm2d (subi vs 1) output_dim in
Layer.of_fn_ (fun xs ~is_training ->
Layer.forward conv xs |> Layer.forward_ bn ~is_training |> relu6)
;;
let inv vs ~stride ~expand_ratio ~input_dim output_dim =
let vs = sub vs "conv" in
let hidden_dim = input_dim * expand_ratio in
let use_residual = input_dim = output_dim && stride = 1 in
let cbr0, nid =
if expand_ratio = 1
then Layer.id_, 0
else cbr (subi vs 0) `k1 ~stride:1 ~input_dim hidden_dim, 1
in
let cbr1 = cbr (subi vs nid) `k3 ~stride ~groups:hidden_dim ~input_dim hidden_dim in
let conv =
Layer.conv2d_
(subi vs (nid + 1))
~ksize:1
~stride:1
~use_bias:false
~input_dim:hidden_dim
output_dim
in
let bn = Layer.batch_norm2d (subi vs (nid + 2)) output_dim in
Layer.of_fn_ (fun xs ~is_training ->
Layer.forward_ cbr0 xs ~is_training
|> Layer.forward_ cbr1 ~is_training
|> Layer.forward conv
|> Layer.forward_ bn ~is_training
|> fun ys -> if use_residual then Tensor.(xs + ys) else ys)
;;
let blocks =
[ 1, 16, 1, 1
; 6, 24, 2, 2
; 6, 32, 3, 2
; 6, 64, 4, 2
; 6, 96, 3, 1
; 6, 160, 3, 2
; 6, 320, 1, 1
]
;;
let v2 vs ~num_classes =
let in_dim = 32 in
let vs_f = sub vs "features" in
let vs_c = sub vs "classifier" in
let init_cbr = cbr (subi vs_f 0) `k3 ~stride:2 ~input_dim:3 in_dim in
let layer_idx = ref 0 in
let last_dim, layers =
List.fold_map blocks ~init:in_dim ~f:(fun in_dim (t, c, nn, s) ->
let layer =
List.range 0 nn
|> List.map ~f:(fun idx ->
Int.incr layer_idx;
let input_dim, stride = if idx = 0 then in_dim, s else c, 1 in
inv (subi vs_f !layer_idx) ~stride ~expand_ratio:t ~input_dim c)
|> Layer.sequential_
in
c, layer)
in
let layers = Layer.sequential_ layers in
Int.incr layer_idx;
let final_cbr = cbr (subi vs_f !layer_idx) `k1 ~stride:1 ~input_dim:in_dim last_dim in
let final_linear = Layer.linear (subi vs_c 1) ~input_dim:last_dim num_classes in
Layer.of_fn_ (fun xs ~is_training ->
Layer.forward_ init_cbr xs ~is_training
|> Layer.forward_ layers ~is_training
|> Layer.forward_ final_cbr ~is_training
|> Tensor.dropout ~p:0.2 ~is_training
|> Tensor.mean_dim ~dim:(Some [ 2 ]) ~keepdim:false ~dtype:(T Float)
|> Tensor.mean_dim ~dim:(Some [ 2 ]) ~keepdim:false ~dtype:(T Float)
|> Layer.forward final_linear)
;;