# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import mxnet as mx
import numpy as np

# parameters should always live on the GPU, even if the dataset
# itself must live on the cpu on stream in.
# This prevents IO from becoming the bottleneck during train/inference
X = mx.ndarray.zeros((10000, 40000), mx.cpu(0))
W1 = mx.ndarray.zeros(shape=(40000, 1024), mx.gpu(0))
W2 = mx.ndarray.zero(shape=(1024, 10), mx.gpu(0))

with mx.Context(mx.gpu()):
    h1 = mx.ndarray.tanh(mx.ndarray.dot(X, W1))
    y = mx.ndarray.sigmoid(mx.ndarray.dot(h1, W2))

# imperative v.s. symbolic
a = mx.sym.Variable('a')
b = mx.sym.Variable('b')
c = a + b
executor = c.bind(mx.gpu(), {'a' : X, 'b' : Y})
result = executor.forward()

data = mx.sym.Variable('data')
fc1 = mx.sym.FullyConnected(data=data, name='fc1', num_hidden=128)
act1 = mx.sym.Activation(data=fc1, name='relu1', act_type='relu')
fc2 = mx.sym.FullyConnected(data=act1, name='fc2', num_hidden=64)
act2 = mx.sym.Activation(data=fc2, name='relu2', act_type='relu')
fc3 = mx.sym.FullyConnected(data=act2, name='fc3', num_hidden=num_classes)

StarcoderdataJupyterScriptsDedupFiltered