#? Lenet5 Convolutional Neural Network Demo
;; simply do {<c> (setq *mnist-dir* "your-mnist-directory")}
;; before loading the present demo. 

(libload "gblearn2/gb-trainers")
(libload "gblearn2/gb-meters")

(libload "gblearn2/net-cscscf")
(libload "gblearn2/demos/dsource-mnist")

(when (not *mnist-dir*)
      (setq *mnist-dir* "./data"))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; a function to load the MNIST database.
(de build-databases-mnist (trsize tesize image-dir)
  (setq trainingdb
        (new dsource-idx3l-narrow
             (new dsource-mnist
                  (load-matrix (concat-fname image-dir "train-images-idx3-ubyte"))
                  (load-matrix (concat-fname image-dir "train-labels-idx1-ubyte"))
		  32 32 0 0.01)
             trsize 0))
  (setq testingdb
        (new dsource-idx3l-narrow
             (new dsource-mnist
                  (load-matrix (concat-fname image-dir "t10k-images-idx3-ubyte"))
                  (load-matrix (concat-fname image-dir "t10k-labels-idx1-ubyte"))
		  32 32 0 0.01)
             tesize (- 5000 (* 0.5 tesize)))) ())


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(de new-lenet5 (image-height 
		image-width
		ki0 kj0 si0 sj0 ki1 kj1 si1 sj1
		hid output-size net-param)
  (let ((table0 (full-table 1 20))
	(table1 (full-table 20 50))
	(table2 (full-table 50 hid)))
    (new net-cscscf
	 image-height image-width
	 ki0 kj0 table0 si0 sj0
	 ki1 kj1 table1 si1 sj1
	 ;; WARNING: those two numbers must be changed 
	 ;; when image-height/image-width change
	 (/ (- (/ (- image-height (1- ki0)) si0) (1- ki1)) si1)
	 (/ (- (/ (- image-width (1- kj0)) sj0) (1- kj1)) sj1)
	 table2
	 output-size
	 net-param)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; number of classes
(setq nclasses 10)
;; the target values for mean-squared error training
;; are +target for the output corresponding to the class 
;; being shown and -target for the other outputs.
(setq target 1)

;; fill matrix with 1-of-n code
(setq labels (int-matrix nclasses))
(setq targets (float-matrix nclasses nclasses))
(idx-f2dotc targets 1.5 targets)
(targets ()() (- target))
(for (i 0 (- nclasses 1)) 
     (targets i i target)
     (labels i i))

;; create trainable parameter
(setq theparam (new idx1-ddparam 0 60000))

;; create the network
(setq thenet
      (new idx3-supervised-module
	   (new-lenet5 32 32 
		       5 5 
		       2 2 
		       5 5 
		       2 2 
		       200  ;; dim of hidden layer
		       10 theparam)
	   (new edist-cost labels 1 1 targets)
	   (new max-classer labels)))

;; create the trainer
(setq thetrainer
      (new supervised-gradient thenet theparam))

;; a classifier-meter measures classification errors
(setq trainmeter (new classifier-meter))
(setq testmeter  (new classifier-meter))

;; initialize the network weights
(==> :thenet:machine forget 1 2)

;; build databases
;; the first two arguments are the sizes of the training
;; and testing sets. The 3rd arg is the directory where
;; the MNIST file reside (change this for your local setup).

(build-databases-mnist 60000 10000 *mnist-dir*)


;; estimate second derivative on 100 iterations, using mu=0.02
;; and set individual espilons
;; (printf "computing diagonal hessian and learning rates\n")
;; (==> thetrainer compute-diaghessian trainingdb 100 0.02)

;; do training iterations 
(printf "training with %d training samples and %d test samples\n" 
	(==> trainingdb size)
	(==> testingdb size))

;; this goes at about 25 examples per second on a PIIIM 800MHz
(de doit (k l)
  (each (((eta n) l))
    (repeat n
      (==> thetrainer compute-diaghessian trainingdb 200 0.02)

      (==> thetrainer train trainingdb trainmeter (* k eta) 0)
      (printf "training: ") (flush)

      (==> thetrainer test trainingdb trainmeter)
      (==> trainmeter display)
      (printf " testing: ") (flush)

      (==> thetrainer test testingdb testmeter)
      (==> testmeter display)  
      ())))

(setq schedule '((1 5) (0.5 4) (0.2 3) (0.1 3)))
(doit 1e-4 schedule)