Jacob's numpy library for machine learning

📰 Attention: You can now purchase jnumpy-themed apparel here.

Getting started

1. Install from pip or clone locally:

$ pip install jnumpy
# or
$ git clone https://github.com/JacobFV/jnumpy.git
$ cd jnumpy
$ pip install .

2. Import the jnumpy module.

importjnumpyasjnp

Examples

Low-level stuff

importjnumpyasjnpW=jnp.Var(np.random.randn(5, 3), trainable=True, name='W')
b_const=jnp.Var(np.array([1., 2., 3.]), name='b') # trainable=False by defaultdefmodel(x):
returnx @ W+b_constdefloss(y, y_pred):
loss= (y-y_pred)**2loss=jnp.ReduceSum(loss, axis=1)
loss=jnp.ReduceSum(loss, axis=0)
returnlossopt=jnp.SGD(0.01)
for_inrange(10):
# make up some datax=jnp.Var(np.random.randn(100, 5))
y=jnp.Var(np.random.randn(100, 3))
# forward passy_pred=model(x)
loss_val=loss(y, y_pred)
# backpropagationopt.minimize(loss)

Neural networks

importjnumpyasjnpimportjnumpy.nnasjnnconv_net=jnn.Sequential(
[
jnn.Conv2D(32, 3, 2, activation=jnp.Relu),
jnn.Conv2D(64, 3, 2, activation=jnp.Relu),
jnn.Flatten(),
jnn.Dense(512, jnp.Sigm),
jnn.Dense(1, jnp.Linear),
]
)

Reinforcement learning

importjnumpyasjnpimportjnumpy.rlasjrlshared_encoder=conv_net# same archiecture as the conv_net above# agentsagentA_hparams= {...}
agentB_hparams= {...}
agentC_hparams= {...}
# categorical deep Q-network:# <q0,q1,..,qn> = dqn(o)# a* = arg_i max qiagentA=jrl.agents.CategoricalDQN(
num_actions=agentA_hparams['num_actions'],
encoder=shared_encoder,
hparams=agentA_hparams,
name='agentA'
)
# standard deep Q-network:# a* = arg_a max dqn(o, a)agentB=jrl.agents.RealDQN(
num_actions=agentB_hparams['num_actions'],
encoder=shared_encoder,
hparams=agentB_hparams,
name='agentB'
)
# random agent:# pick a random actionagentC=jrl.agents.RandomAgent(agentC_hparams['num_actions'], name='agentC')
# init enviromentstrain_env=jrl.ParallelEnv(
batch_size=32,
env_init_fn=lambda: MyEnv(...), # `jrl.Environment` subclass. Must have `reset` and `step` methods.
)
dev_env=jrl.ParallelEnv(
batch_size=8,
env_init_fn=lambda: MyEnv(...),
)
test_env=jrl.ParallelEnv(
batch_size=8,
env_init_fn=lambda: MyEnv(...),
)
# traintrainer=jrl.ParallelTrainer(callbacks=[
jrl.PrintCallback(['epoch', 'agent', 'collect_reward', 'q_train', 'q_test']),
jrl.QEvalCallback(eval_on_train=True, eval_on_test=True),
])
trainer.train(
agents={'agentA': agentA, 'agentB': agentB},
all_hparams={'agentA': agentA_hparams, 'agentB': agentB_hparams},
env=train_env,
test_env=dev_env,
training_epochs=10,
)
# testdriver=ParallelDriver()
trajs=driver.drive(
agents={'agentA': agentA, 'agentB': agentB},
env=test_env
)
per_agent_rewards= {
agent_name: sum(step.rewardforstepintraj)
foragent_name, trajintrajs.items()}
print('cumulative test rewards:', per_agent_rewards)

Limitations and Future Work

Future versions will feature:

• add fit, evaluate, and predict to jnp.Sequential
• recurrent network layers
• static execution graphs allowing breadth-first graph traversal
• more optimizers, metrics, and losses
• io loaders for csv's, images, and models (maybe also for graphs)
• more examples

Also maybe for the future:

• custom backends (i.e.: tensorflow or pytorch instead of numpy)

License

All code in this repository is licensed under the MIT license. No restrictions, but no warranties. See the LICENSE file for details.

Contributing

This is a small project, and I don't plan on growing it much. You are welcome to fork and contribute or email me jacob [dot] valdez [at] limboid [dot] ai if you would like to take over. You can add your name to the copyright if you make a PR or your own branch.

The codebase is kept in only a few files, and I have tried to minimize the use of module prefixes because my CSE 4308/4309/4392 classes require the submissions to be stitched togethor in a single file.