clojure-tensorflow/src/clojure_tensorflow/utils.clj at abd7903b563d08051ac1608185e0d4aa314a21bb · kieranbrowne/clojure-tensorflow · GitHub

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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
(ns clojure-tensorflow.utils
  (:require [clojure-tensorflow.graph :as graph]
            [clojure-tensorflow.session :as session])
  (:import [org.tensorflow
            Tensor
            Shape
            DataType]))

(defn recursively
  "Apply function to all items in nested data structure if
  condition function is met."
  [apply-if-fn func data]
  (if (apply-if-fn data)
    (func (map (partial recursively apply-if-fn func) data))
    data))

(defn make-coll
  "Make a collection of x,y,z... dimensions"
  [fill & dims]
  (case (count dims)
    0 fill
    1 (repeat (first dims) fill)
    (repeat (first dims) (apply (partial make-coll fill) (rest dims)))
    ))

(def array?
  "Works like coll? but returns true if argument is array"
  #(= \[ (first (.getName (.getClass %)))))

(def tensor->shape
  #(let [arr (.shape %)]
     (if (> (count arr) 0)
       (Shape/make
        (aget arr 0)
        (java.util.Arrays/copyOfRange arr 1 (count arr)))
       (Shape/scalar))))

(defn tf-vals [v]
  "Convert value into type acceptable to TensorFlow
  Persistent data structures become arrays
  Longs become 32bit integers
  Doubles become floats"
  (cond
    (coll? v)
    (if (coll? (first v))
      (into-array (map tf-vals v))
      (case (.getName (type (first v)))
        "java.lang.Long" (int-array v)
        "java.lang.Int" (int-array v)
        "java.lang.Double" (float-array v)
        "java.lang.Float" (float-array v)))
    (= (.getName (type v)) "java.lang.Long") (int v)
    (= (.getName (type v)) "java.lang.Int") (int v)
    (= (.getName (type v)) "java.lang.Double") (float v)
    (= (.getName (type v)) "java.lang.Float") (float v)
    ;; anything else
    true v))


(defn output-shape [tensor]
  (let [shape (tensor->shape tensor)
        dims (map #(.size shape %)
                  (range (.numDimensions shape)))
        d (case (.name (.dataType tensor))
            "INT32" 0
            "INT64" 0
            "FLOAT" 0.0
            "DOUBLE" 0.0)]
    (tf-vals
     (apply (partial make-coll d) dims))))

(defn get-tensor-val [tensor]
  (let [copy-to (output-shape tensor)
        dtype (.name (.dataType tensor))
        ]
    (cond
      (array? copy-to) (.copyTo tensor copy-to)
      (= "FLOAT" dtype) (double (.floatValue tensor))
      (= "DOUBLE" dtype) (.doubleValue tensor)
      (= "INT32" dtype) (long (.intValue tensor))
      (= "INT64" dtype) (.longValue tensor))))


(def tensor->clj (comp (partial recursively array? vec) get-tensor-val))

(def clj->tensor #(Tensor/create (tf-vals %)))

(defn tf-obj? [x]
  (if (re-find #"org.tensorflow" (.getName (class x)))
    true false))

(def to-floats (partial clojure.walk/postwalk #(if (coll? %) % (float %))))

(defn thread
  "Approximately equivalent to -> macro.
  Required because -> must run at compile time"
  [val functions] (reduce #(%2 %1) val functions))

(defprotocol Conversion
  (->clj [x] "Return clojure values")
  (->tensor [x] "Return tensor values")
  (->output [x] "Return tensorflow output object")
  )

(extend-type org.tensorflow.Shape
  Conversion
  (->clj [x] (map (fn [i] (.size x i))
                  (range (.numDimensions x)))))

(extend-type org.tensorflow.Tensor
  Conversion
  (->output [x]
    (-> graph/graph
        (.opBuilder "Const" (str (gensym)))
        (.setAttr "dtype" (.dataType x))
        (.setAttr "value" x)
        .build
        (.output 0)))

  (->clj [x]
    ((comp (partial recursively array? vec) get-tensor-val) x))
  )

(extend-type org.tensorflow.Output
  Conversion
  (->tensor [x]
    (-> session/session .runner (.fetch (.name (.op x))) .run (.get 0))))

(extend-protocol Conversion
  java.lang.Long
  (->tensor [x] (->tensor (int x)))
  java.lang.Double
  (->tensor [x] (->tensor (float x)))
  java.lang.Float
  (->tensor [x] (Tensor/create x))
  java.lang.Integer
  (->tensor [x] (Tensor/create x))
  clojure.lang.PersistentVector
  (->tensor [x] (Tensor/create (tf-vals x)))
  clojure.lang.PersistentList
  (->tensor [x] (->tensor (vec x))))