def test_cloudpickle(self):
model = eg.Model(
module=eg.Linear(10),
loss=[
eg.losses.Crossentropy(),
eg.regularizers.L2(1e-4),
],
metrics=eg.metrics.Accuracy(),
optimizer=optax.adamw(1e-3),
eager=True,
)
X = np.random.uniform(size=(5, 2))
y = np.random.randint(10, size=(5, ))
y0 = model.predict(X)
with TemporaryDirectory() as model_dir:
model.save(model_dir)
newmodel = eg.load(model_dir)
y1 = newmodel.predict(X)
assert np.all(y0 == y1)
def main(
debug: bool = False,
eager: bool = False,
logdir: str = "runs",
steps_per_epoch: int = 200,
epochs: int = 100,
batch_size: int = 64,
):
if debug:
import debugpy
print("Waiting for debugger...")
debugpy.listen(5678)
debugpy.wait_for_client()
current_time = datetime.now().strftime("%b%d_%H-%M-%S")
logdir = os.path.join(logdir, current_time)
X_train, y_train, X_test, y_test = dataget.image.mnist(
global_cache=True).get()
X_train = X_train[..., None]
X_test = X_test[..., None]
print("X_train:", X_train.shape, X_train.dtype)
print("y_train:", y_train.shape, y_train.dtype)
print("X_test:", X_test.shape, X_test.dtype)
print("y_test:", y_test.shape, y_test.dtype)
class CNN(elegy.Module):
def call(self, image: jnp.ndarray, training: bool):
@elegy.to_module
def ConvBlock(x, units, kernel, stride=1):
x = elegy.nn.Conv2D(units,
kernel,
stride=stride,
padding="same")(x)
x = elegy.nn.BatchNormalization()(x, training)
x = elegy.nn.Dropout(0.2)(x, training)
return jax.nn.relu(x)
x: np.ndarray = image.astype(jnp.float32) / 255.0
# base
x = ConvBlock()(x, 32, [3, 3])
x = ConvBlock()(x, 64, [3, 3], stride=2)
x = ConvBlock()(x, 64, [3, 3], stride=2)
x = ConvBlock()(x, 128, [3, 3], stride=2)
# GlobalAveragePooling2D
x = jnp.mean(x, axis=[1, 2])
# 1x1 Conv
x = elegy.nn.Linear(10)(x)
return x
model = elegy.Model(
module=CNN(),
loss=elegy.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=elegy.metrics.SparseCategoricalAccuracy(),
optimizer=optax.adam(1e-3),
run_eagerly=eager,
)
# show model summary
model.summary(X_train[:64], depth=1)
history = model.fit(
x=X_train,
y=y_train,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
batch_size=batch_size,
validation_data=(X_test, y_test),
shuffle=True,
callbacks=[TensorBoard(logdir=logdir)],
)
elegy.utils.plot_history(history)
model.save("models/conv")
model = elegy.load("models/conv")
print(model.evaluate(x=X_test, y=y_test))
# get random samples
idxs = np.random.randint(0, 10000, size=(9, ))
x_sample = X_test[idxs]
# get predictions
y_pred = model.predict(x=x_sample)
# plot results
with SummaryWriter(os.path.join(logdir, "val")) as tbwriter:
figure = plt.figure(figsize=(12, 12))
for i in range(3):
for j in range(3):
k = 3 * i + j
plt.subplot(3, 3, k + 1)
plt.title(f"{np.argmax(y_pred[k])}")
plt.imshow(x_sample[k], cmap="gray")
# tbwriter.add_figure("Conv classifier", figure, 100)
plt.show()
def main(
debug: bool = False,
eager: bool = False,
logdir: str = "runs",
steps_per_epoch: int = 200,
epochs: int = 100,
batch_size: int = 64,
):
if debug:
import debugpy
print("Waiting for debugger...")
debugpy.listen(5678)
debugpy.wait_for_client()
current_time = datetime.now().strftime("%b%d_%H-%M-%S")
logdir = os.path.join(logdir, current_time)
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
def preprocess_images(images):
images = images.reshape((images.shape[0], 28, 28, 1)) / 255.0
return images.astype("float32")
X_train = preprocess_images(X_train)
X_test = preprocess_images(X_test)
print("X_train:", X_train.shape, X_train.dtype)
print("y_train:", y_train.shape, y_train.dtype)
print("X_test:", X_test.shape, X_test.dtype)
print("y_test:", y_test.shape, y_test.dtype)
model = eg.Model(
module=CNN(),
loss=eg.losses.Crossentropy(),
metrics=eg.metrics.Accuracy(),
optimizer=optax.adam(1e-3),
eager=eager,
)
# show summary
model.summary(X_train[:64])
batch_size = 64
train_size = 60000
test_size = 10000
# Create tf datasets
train_dataset = (tf.data.Dataset.from_tensor_slices(
(X_train, y_train)).shuffle(train_size).batch(batch_size).repeat())
test_dataset = (tf.data.Dataset.from_tensor_slices(
(X_test, y_test)).shuffle(test_size).batch(batch_size))
history = model.fit(
train_dataset,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
validation_data=test_dataset,
callbacks=[eg.callbacks.TensorBoard(logdir=logdir)],
)
eg.utils.plot_history(history)
model.save("models/conv")
model = eg.load("models/conv")
print(model.evaluate(x=X_test, y=y_test))
# get random samples
idxs = np.random.randint(0, 10000, size=(9, ))
x_sample = X_test[idxs]
# get predictions
y_pred = model.predict(x=x_sample)
# plot results
figure = plt.figure(figsize=(12, 12))
for i in range(3):
for j in range(3):
k = 3 * i + j
plt.subplot(3, 3, k + 1)
plt.title(f"{np.argmax(y_pred[k])}")
plt.imshow(x_sample[k], cmap="gray")
plt.show()
def main(
debug: bool = False,
eager: bool = False,
logdir: str = "runs",
steps_per_epoch: int = 200,
epochs: int = 100,
batch_size: int = 64,
):
if debug:
import debugpy
print("Waiting for debugger...")
debugpy.listen(5678)
debugpy.wait_for_client()
current_time = datetime.now().strftime("%b%d_%H-%M-%S")
logdir = os.path.join(logdir, current_time)
dataset = load_dataset("mnist")
dataset.set_format("np")
X_train = np.stack(dataset["train"]["image"])[..., None]
y_train = dataset["train"]["label"]
X_test = np.stack(dataset["test"]["image"])[..., None]
y_test = dataset["test"]["label"]
print("X_train:", X_train.shape, X_train.dtype)
print("y_train:", y_train.shape, y_train.dtype)
print("X_test:", X_test.shape, X_test.dtype)
print("y_test:", y_test.shape, y_test.dtype)
model = eg.Model(
module=CNN(),
loss=eg.losses.Crossentropy(),
metrics=eg.metrics.Accuracy(),
optimizer=optax.adam(1e-3),
eager=eager,
)
# show summary
model.summary(X_train[:64])
train_dataset = TensorDataset(torch.from_numpy(X_train),
torch.from_numpy(y_train))
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
test_dataset = TensorDataset(torch.from_numpy(X_test),
torch.from_numpy(y_test))
test_dataloader = DataLoader(test_dataset, batch_size=batch_size)
history = model.fit(
train_dataloader,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
validation_data=test_dataloader,
callbacks=[eg.callbacks.TensorBoard(logdir=logdir)],
)
eg.utils.plot_history(history)
model.save("models/conv")
model = eg.load("models/conv")
print(model.evaluate(x=X_test, y=y_test))
# get random samples
idxs = np.random.randint(0, 10000, size=(9, ))
x_sample = X_test[idxs]
# get predictions
y_pred = model.predict(x=x_sample)
# plot results
with SummaryWriter(os.path.join(logdir, "val")) as tbwriter:
figure = plt.figure(figsize=(12, 12))
for i in range(3):
for j in range(3):
k = 3 * i + j
plt.subplot(3, 3, k + 1)
plt.title(f"{np.argmax(y_pred[k])}")
plt.imshow(x_sample[k], cmap="gray")
# tbwriter.add_figure("Conv classifier", figure, 100)
plt.show()