# Training and CV sets
trX = tr1X[:50000]
trY = tr0Y[:50000]
cvX = tr1X[50000:]
cvY = tr0Y[50000:]
# RBMs
rbmobject1 = RBM(FLAGS.out_dir, 784, 900, ['rbmw1', 'rbvb1', 'rbmhb1'], 0.3,
tf.nn.sigmoid)
rbmobject2 = RBM(FLAGS.out_dir, 900, 500, ['rbmw2', 'rbvb2', 'rbmhb2'], 0.3,
tf.nn.sigmoid)
rbmobject1.restore_weights('./out/rbmw1.chp')
rbmobject2.restore_weights('./out/rbmw2.chp')
# DBN
y_colNb = 1
dbn = DBN(FLAGS.out_dir, [rbmobject1, rbmobject2], 784, y_colNb,
FLAGS.class_nb, FLAGS.batchsize, FLAGS.epochs, 0.3)
#dbn.restore_trained_net('./out/dbn.chp')
#i=8
#dbn.predict(teX[i], teY[i])
#exit()
iterations = len(trX) / FLAGS.batchsize
## Train DBN
print 'Training DBN: ', FLAGS.epochs, ' epochs of ', iterations, ' iterations'
dbn.train(trX, trY, cvX, cvY, teX, teY)
dbn.save_trained_net('./out/dbn.chp')
