cgarciae

MyModule:
  b: list
    - Initializer
    - Parameter((5, 13), float32)
  a: dict
    mlps: list
      - MLP:
          linear1: Linear
            w: Parameter((2, 3), float32)
            b: Parameter((3,), float32)

cgarciae

# PARAMETERS
MODEL = "ResNet50"
OUTPUT_DIRECTORY = "models/resnet50"
EPOCHS = 90
BATCH_SIZE = 6
IMAGE_SIZE = 224  # image size in pixels
DATASET = "imagenet2012:5.1.*"  # TFDS dataset name and version
DTYPE = "float16"  # float16 for mixed_precision or float32 for normal mode
LEARNING_RATE = 0.1 * BATCH_SIZE / 256.0
MOMENTUM = 0.9

cgarciae

model = elegy.Model(
    module=MixtureModel(k=k), 
    loss=MixtureNLL(), 
    optimizer=optax.adam(3e-4),
)

model.summary(X_train[:batch_size], depth=1)

model.fit(
    x=X_train, 

cgarciae

class MixtureNLL(elegy.Loss):
    def call(self, y_true, y_pred):
        y, probs = y_pred
        y_true = jnp.broadcast_to(y_true, (y_true.shape[0], y.shape[1]))

        return -safe_log(
            jnp.sum(
                probs
                * jax.scipy.stats.norm.pdf(y_true, loc=y[..., 0], scale=y[..., 1]),
                axis=1,

cgarciae

class MixtureModel(elegy.Module):
    def __init__(self, k: int):
        super().__init__()
        self.k = k

    def call(self, x):

        x = elegy.nn.Linear(64, name="backbone")(x)
        x = jax.nn.relu(x)

cgarciae

import numpy as np
import jax
import jax.numpy as jnp
import elegy
import optax

cgarciae

import numpy as np
import jax
import jax.numpy as jnp
import elegy
import optax

cgarciae

class Accuracy(hk.Module):
    def __init__(self, y_true: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray:

        total = hk.get_state(
            "total", shape=[], dtype=jnp.float32, init=hk.initializers.Constant(0)
        )
        count = hk.get_state(
            "count", shape=[], dtype=jnp.int64, init=hk.initializers.Constant(0)
        )

cgarciae

from aiohttp import ClientSession, TCPConnector
import asyncio
import sys
import pypeln as pl


limit = 1000
urls = ("http://localhost:8080/{}".format(i) for i in range(int(sys.argv[1])))

cgarciae

import pypeln as pl
import time
from random import random

def slow_add1(x):
  time.sleep(random()) # <= some slow computation
  return x + 1

def slow_gt3(x):
  time.sleep(random()) # <= some slow computation