!pip install tensorflow-gpu

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0,1

import os
print(os.environ["CUDA_DEVICE_ORDER"])
print(os.environ["CUDA_VISIBLE_DEVICES"])

import tensorflow as tf

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

tf.test.gpu_device_name()

"""
Explanation: It works on a Windows 10 PC which has an old GPU: GeForce GTX 550 Ti.
Simply run the following command (not using docker) in the root of this project:
jupyter notebook
Also run the following in a separate terminal:
tensorboard --logdir=./dl/logs --host=127.0.0.1 --port=8889
End of explanation
"""


!nvcc --version

!pip install mxnet-cu90

# From https://stackoverflow.com/questions/49076092/is-there-a-way-to-check-if-mxnet-uses-my-gpu/49079940#49079940
# https://developer.download.nvidia.com/compute/DCGM/docs/nvidia-smi-367.38.pdf
!"C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi" --query-gpu=timestamp,name,pci.bus_id,driver_version,pstate,pcie.link.gen.max,pcie.link.gen.current,temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.free,memory.used --format=csv

import mxnet as mx 
def gpu_device(gpu_number=0):
    try:
        _ = mx.nd.array([1, 2, 3], ctx=mx.gpu(gpu_number))
    except mx.MXNetError:
        return None
    return mx.gpu(gpu_number)

gpu_device()

mx.gpu(0)

"""
Explanation: From https://mxnet.incubator.apache.org/install/index.html :
End of explanation
"""


from __future__ import print_function
import numpy as np
import mxnet as mx
from mxnet import nd, autograd, gluon

data_ctx = mx.cpu()
model_ctx = mx.cpu()
# model_ctx = mx.gpu(0)

batch_size = 64
num_inputs = 784
num_outputs = 10
num_examples = 60000
def transform(data, label):
    return data.astype(np.float32)/255, label.astype(np.float32)

train_data = mx.gluon.data.DataLoader(mx.gluon.data.vision.MNIST(train=True, transform=transform),
                                      batch_size, shuffle=True)
test_data = mx.gluon.data.DataLoader(mx.gluon.data.vision.MNIST(train=False, transform=transform),
                                     batch_size, shuffle=False)

class MLP(gluon.Block):
    def __init__(self, **kwargs):
        super(MLP, self).__init__(**kwargs)
        with self.name_scope():
            self.dense0 = gluon.nn.Dense(64)
            self.dense1 = gluon.nn.Dense(64)
            self.dense2 = gluon.nn.Dense(10)

    def forward(self, x):
        x = nd.relu(self.dense0(x))
        x = nd.relu(self.dense1(x))
        x = self.dense2(x)
        return x

net = MLP()
net.collect_params().initialize(mx.init.Normal(sigma=.01), ctx=model_ctx)

data = nd.ones((1,784))
net(data.as_in_context(model_ctx))

class MLP(gluon.Block):
    def __init__(self, **kwargs):
        super(MLP, self).__init__(**kwargs)
        with self.name_scope():
            self.dense0 = gluon.nn.Dense(64, activation="relu")
            self.dense1 = gluon.nn.Dense(64, activation="relu")
            self.dense2 = gluon.nn.Dense(10)

    def forward(self, x):
        x = self.dense0(x)
        print("Hidden Representation 1: %s" % x)
        x = self.dense1(x)
        print("Hidden Representation 2: %s" % x)
        x = self.dense2(x)
        print("Network output: %s" % x)
        return x

net = MLP()
net.collect_params().initialize(mx.init.Normal(sigma=.01), ctx=model_ctx)
net(data.as_in_context(model_ctx))

num_hidden = 64
net = gluon.nn.HybridSequential()
with net.name_scope():
    net.add(gluon.nn.Dense(num_hidden, activation="relu"))
    net.add(gluon.nn.Dense(num_hidden, activation="relu"))
    net.add(gluon.nn.Dense(num_outputs))
net.hybridize()

net.collect_params().initialize(mx.init.Normal(sigma=.1), ctx=model_ctx)

softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()

trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': .01})

def evaluate_accuracy(data_iterator, net):
    acc = mx.metric.Accuracy()
    for i, (data, label) in enumerate(data_iterator):
        data = data.as_in_context(model_ctx).reshape((-1, 784))
        label = label.as_in_context(model_ctx)
        output = net(data)
        predictions = nd.argmax(output, axis=1)
        acc.update(preds=predictions, labels=label)
    return acc.get()[1]

!pip install mxboard

!pip install tensorboard

from mxboard import SummaryWriter
sw = SummaryWriter(logdir='logs', flush_secs=5)

epochs = 10
smoothing_constant = .01

# collect parameter names for logging the gradients of parameters in each epoch
params = net.collect_params()
param_names = params.keys()
global_step = 0

for e in range(epochs):
    cumulative_loss = 0
    for i, (data, label) in enumerate(train_data):
        data = data.as_in_context(model_ctx).reshape((-1, 784))
        label = label.as_in_context(model_ctx)
        with autograd.record():
            output = net(data)
            loss = softmax_cross_entropy(output, label)
 
        sw.add_scalar(tag='cross_entropy', value=loss.mean().asscalar(), global_step=global_step)
        if i == 0:
            sw.add_image('minist_first_minibatch', data.reshape((batch_size, 1, 28, 28)), e)
        if e == 0:
            sw.add_graph(net)
        grads = [i.grad() for i in net.collect_params().values()]
        for i, name in enumerate(param_names):
            sw.add_histogram(tag=name, values=grads[i], global_step=e, bins=1000)

        global_step += 1
        loss.backward()
        trainer.step(data.shape[0])
        cumulative_loss += nd.sum(loss).asscalar()

        
    
    test_accuracy = evaluate_accuracy(test_data, net)
    train_accuracy = evaluate_accuracy(train_data, net)
    sw.add_scalar(tag='accuracy_curves', value=('train_acc', train_accuracy), global_step=e)
    sw.add_scalar(tag='accuracy_curves', value=('valid_acc', test_accuracy), global_step=e)
    print("Epoch %s. Loss: %s, Train_acc %s, Test_acc %s" %
          (e, cumulative_loss/num_examples, train_accuracy, test_accuracy))

"""
Explanation: From https://gluon.mxnet.io/chapter03_deep-neural-networks/mlp-gluon.html :
End of explanation
"""


