!pip install mxnet
!pip install opencv-python
!pip install matplotlib

import mxnet as mx
import matplotlib.pyplot as plt
import numpy as np
from collections import namedtuple

#test that MXnet works on hosted runtime (I said this to make me sound smart)

test_images = []
for i in range(1,10):
    x_data = np.random.randint(0,255,[3,256,256])
    test_images.append(x_data)

test_labels = np.ones(10)

'''
#Define iterator class
class RandIter(mx.io.DataIter):

    def __init__ (self,batch_size,ndim):
        self.batch_size = batch_size
        self.ndim = ndim
        self.provide_data = [('data'),(batch_size,ndim,256,256)]
        self.provide_label = []

    def iter_next(self):
        return True

    def getdata(self):
        return [mx.random.normal(0,1,shape=(self.batch_size,self.ndim,256,256))]

x = RandIter(1,3)
'''

Batch = namedtuple('Batch', ['data'])

test_batch = Batch(test)

test = mx.io.NDArrayIter(test_images,label=test_labels,batch_size=10)

test_batch = Batch(test)
#rand_iter = RandIter(batch_size=10,ndim=3)

#Define Network Architecture
data = mx.sym.Variable('data')
g1 = mx.sym.Convolution(data, name = 'CONV1',kernel=(1,1),num_filter=1000)
g2 = mx.sym.Convolution(g1,name = 'Conv2', kernel=(1,1),num_filter=3)

conv_model = mx.mod.Module(symbol=g2,data_names=('data',),label_names=None,context=mx.cpu())

conv_model.bind(for_training=True,data_shapes=[('data',(1,3,256,256))])
conv_model.init_params(initializer=mx.init.Xavier(magnitude=2.))
conv_model.init_optimizer(
            optimizer='adam',
            optimizer_params= {'learning_rate': 0.0002,'beta1':0.5})

conv_model.forward(test_batch)
z = conv_model.get_outputs()

x = z[0].asnumpy()
x = x.reshape(3,256,256)

#plt.imshow(x)