Merge pull request #1 from taraspiotr/tensor

tensor

src/DeepJulia.jl

@@ -3,8 +3,12 @@ module DeepJulia
 export Device,
 gpu,
 cpu,
-ArrayOrCuArray,
+
+# tensor
+
+Tensor,
 to,
+device,
 
 # variable
 
@@ -44,7 +48,8 @@ batchify,
 FashionMNIST,
 shuffle!
 
-include("cuda.jl")
+include("device.jl")
+include("tensor.jl")
 include("variable.jl")
 include("loss.jl")
 include("modules.jl")

src/device.jl

@@ -0,0 +1,4 @@
+@enum Device begin
+    cpu
+    gpu
+end

src/modules.jl

@@ -12,9 +12,9 @@ mutable struct LinearLayer <: NNModule
     W::Variable
     b::Variable
 
-    input::ArrayOrCuArray
+    input::Tensor
 
-    LinearLayer(W::ArrayOrCuArray, b::ArrayOrCuArray) = new(
+    LinearLayer(W::Tensor, b::Tensor) = new(
         Variable(W),
         Variable(b),
         Matrix{Real}(undef, 0, 0),

src/optim.jl

@@ -7,7 +7,7 @@ struct SGD <: Optimizer
     params::Vector{Variable}
     lr::Real
     momentum::Real
-    velocities::Vector{ArrayOrCuArray}
+    velocities::Vector{Tensor}
 
     SGD(params, lr) = new(params, lr, 0.0, Vector{Matrix{Real}}(undef, size(params, 1)))
     SGD(params, lr, momentum) = new(params, lr, momentum, [to(zeros(size(p.values)), device(p.values)) for p ∈ params])

src/cuda.jl → src/tensor.jl

@@ -1,13 +1,8 @@
 using CUDA: CuArray
 
-@enum Device begin
-    cpu
-    gpu
-end
-
-ArrayOrCuArray = Union{Array,CuArray}
+Tensor = Union{Array,CuArray}
 
-function to(A::ArrayOrCuArray, device::Device)
+function to(A::Tensor, device::Device)
     if device == cpu
         return Array(A)
     elseif device == gpu
@@ -17,4 +12,4 @@ function to(A::ArrayOrCuArray, device::Device)
     end
 end
 
-device(A::ArrayOrCuArray) = isa(A, CuArray) ? gpu : cpu
+device(A::Tensor) = isa(A, CuArray) ? gpu : cpu

src/variable.jl

@@ -1,6 +1,6 @@
 mutable struct Variable
-    values::ArrayOrCuArray
-    grad::ArrayOrCuArray
+    values::Tensor
+    grad::Tensor
 
     Variable(values) = new(values, to(zeros(size(values)), device(values)))
     Variable(values, grad) = new(values, grad)