scripts/nn/examples/mnist_resnet.dml - systemds - Git at Google

 #-------------------------------------------------------------
 #
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
 #
 #-------------------------------------------------------------

 # MNIST Resnet - Train
 #
 # This script trains a convolutional net using the "ResNet" architecture
 # on images of handwritten digits.
 #
 # Inputs:
 #  - train: File containing labeled MNIST training images.
 #     The format is "label, pixel_1, pixel_2, ..., pixel_n".
 #  - test: File containing labeled MNIST test images.
 #     The format is "label, pixel_1, pixel_2, ..., pixel_n".
 #  - C: Number of color chanels in the images.
 #  - Hin: Input image height.
 #  - Win: Input image width.
 #  - epochs: [DEFAULT: 10] Total number of full training loops over
 #     the full data set.
 #  - out_dir: [DEFAULT: "."] Directory to store weights and bias
 #     matrices of trained model, as well as final test accuracy.
 #  - fmt: [DEFAULT: "csv"] File format of `train` and `test` data.
 #     Options include: "csv", "mm", "text", and "binary".
 #
 # Outputs:
 #  - W1, W2, W3, W4: Files containing the trained weights of the model.
 #  - b1, b2, b3, b4: Files containing the trained biases of the model.
 #  - accuracy: File containing the accuracy and loss on the test data over all epochs.
 #
 # Data:
 # The MNIST dataset contains labeled images of handwritten digits,
 # where each example is a 28x28 pixel image of grayscale values in
 # the range [0,255] stretched out as 784 pixels, and each label is
 # one of 10 possible digits in [0,9].
 #
 # Sample Invocation (running from outside the `nn` folder):
 # 1. Download data (60,000 training examples, and 10,000 test examples)
 #   ```
 #   nn/examples/get_mnist_data.sh
 #   ```
 #
 # 2. Execute using Spark
 #   ```
 #   spark-submit --master local[*] --driver-memory 10G
 #   --conf spark.driver.maxResultSize=0 --conf spark.rpc.message.maxSize=128
 #   $SYSTEMDS_ROOT/target/SystemDS.jar -f nn/examples/mnist_resnet.dml
 #   -nvargs train=nn/examples/data/mnist/mnist_train.csv test=nn/examples/data/mnist/mnist_test.csv
 #   C=1 Hin=28 Win=28 epochs=10 out_dir=nn/examples/model/mnist_resnet
 #   ```
 #

 source("nn/networks/resnet18.dml") as resnet
 source("scripts/nn/layers/softmax_cross_entropy_loss.dml") as loss_nn

 # Read the data
 fmt = "csv"
 train = read("scripts/nn/examples/data/mnist_train.csv", format=fmt)
 test = read("scripts/nn/examples/data/mnist_test.csv", format=fmt)
 out_dir = "scripts/nn/example/model/mnist_resnet"

 # Extract images and labels
 images = train[,2:ncol(train)]
 labels = train[,1]
 images_test = test[,2:ncol(test)]
 labels_test = test[,1]
 classes = 10

 # Scale images to [-1,1], and one-hot encode the labels
 N = nrow(images)
 N_test = nrow(images_test)
 X = (images / 255.0) * 2 - 1
 X = cbind(X, X, X) # Resnet assumes C=3 so we duplicate the data along the channels
 Y = table(seq(1, N), labels+1, N, 10)
 X_test = (images_test / 255.0) * 2 - 1
 X_test = cbind(X_test, X_test, X_test)
 Y_test = table(seq(1, N_test), labels_test+1, N_test, 10)

 # Split into training (55,000 examples) and validation (5,000 examples)
 #X = images[5001:nrow(images),]
 #X_val = images[1:5000,]
 #Y = labels[5001:nrow(images),]
 #Y_val = labels[1:5000,]

 # Get initial model parameters
 [model, ema_means_vars] = resnet::init(classes, -1)

 # Get initial optimizer parameters
 optimizer_params = resnet::init_lars_optim_params(classes)
 # optimizer_params = resnet::init_sgd_momentum_optim_params(classes)
 # optimizer_params = resnet::init_adam_optim_params(classes)

 # Define image properties
 Hin = 28
 Win = 28
 #N_val = 0

 # Define training parameters
 epochs = 90
 batch_size = 512

 [accuracy, loss_metric, learned_model, learned_emas] = train(X, Y, X_test, Y_test, model, ema_means_vars, N, Hin, Win, epochs, batch_size, optimizer_params)

 write(accuracy, "scripts/nn/examples/out/resnet_mnist_accuracy.csv", format="csv")
 write(loss_metric, "scripts/nn/examples/out/resnet_mnist_loss.csv", format="csv")

 #Train
 train = function(matrix[double] X, matrix[double] Y, matrix[double] X_test, matrix[double] Y_test, list[unknown] model, list[unknown] emas, int samples, int Hin,
     int Win, int epochs, int batch_size, list[unknown] optim_params)
     return (matrix[double] accuracy, matrix[double] loss_metric,
             list[unknown] learned_model, list[unknown] learned_emas) {

     # --- LEARNING RATE SCHEDULE HYPERPARAMETERS ---
     # The learning rate we want to reach AFTER warmup
     initial_lr = 0.01
     # A very small final learning rate to decay towards
     end_lr = 0.0001
     # Number of warmup epochs, as per the paper
     warmup_epochs = 5
     # The exponent for the polynomial decay, as per the paper
     power = 2.0

     # Optimizer hyperparameters
     momentum = 0.9
     trust_coeff = 0.001
     weight_decay = 0.0001

     # Adam optimizer hyperparameters
     beta1 = 0.9
     beta2 = 0.999
     epsilon = 1e-8

     # Calculate total iterations for the schedule
     iterations_per_epoch = ceil(samples / batch_size)
     total_iterations = epochs * iterations_per_epoch
     warmup_iterations = warmup_epochs * iterations_per_epoch
     decay_iterations = total_iterations - warmup_iterations

     # Initialize metrics
     learned_model = list()
     learned_emas = list()
     accuracy = matrix(0, rows=epochs, cols=1)
     loss_metric = matrix(0, rows=epochs, cols=1)

     iterations = ceil(samples/batch_size)
     mode = "train"

     for (epoch in 1:epochs) {
         loss_avg = 0.0

         print("Start epoch: " + epoch + "/" + epochs)

         for (i in 1:iterations) {
             print(" - Iteration: " + i + "/" + iterations)

             # --- START DYNAMIC LEARNING RATE LOGIC ---
             current_iteration = (epoch - 1) * iterations_per_epoch + i
             current_lr = 0.0

             if (current_iteration < warmup_iterations) {
                 # 1. Linear Warmup Phase
                 # Linearly increase LR from 0 to initial_lr over warmup_iterations
                 current_lr = initial_lr * (as.double(current_iteration) / warmup_iterations)
             } else {
                 # 2. Polynomial Decay Phase
                 decay_step = current_iteration - warmup_iterations
                 decay_progress = as.double(decay_step) / decay_iterations
                 current_lr = end_lr + (initial_lr - end_lr) * (1 - decay_progress)^power
             }

             if (i == 1) { # Print LR once per epoch to reduce log spam
                 print("Using Learning Rate: " + current_lr)
             }
             # --- END DYNAMIC LEARNING RATE LOGIC ---

             # Get batch
             start = (i - 1) * batch_size + 1
             end = min(samples, i * batch_size)
             X_batch = X[start:end,]
             Y_batch = Y[start:end,]

             # Forward pass
             [out, emas, cached_out, cached_means_vars] = resnet::forward(X_batch, Hin, Win, model, mode, emas)

             # Loss
             loss = loss_nn::forward(out, Y_batch)
             if (i %% 10 == 0) { # Print loss less frequently on large datasets
                 print(" - Iteration: " + i + "/" + iterations + ", Loss: " + loss)
             }
             loss_avg = (loss_avg * (i - 1) + loss) / i

             # Backward
             dOut = loss_nn::backward(out, Y_batch)
             [dX, gradients] = resnet::backward(dOut, cached_out, model, cached_means_vars)

             # Update parameters
             [model, optim_params] = resnet::update_params_with_lars(model, gradients, current_lr, momentum, weight_decay, trust_coeff,
                   optim_params)
             # [model, optim_params] = resnet::update_params_with_sgd_momentum(model, gradients, current_lr, momentum, optim_params)

             # [model, optim_params] = resnet::update_params_with_adam(model, gradients, current_lr, beta1, beta2, epsilon, current_iteration, optim_params)
         }

         # Reshuffle mini batches
         r = rand(rows=nrow(Y), cols=1, min=0, max=1, pdf="uniform")
         X_tmp = order(target=cbind(r, X), by=1)
         Y_tmp = order(target=cbind(r, Y), by=1)
         X = X_tmp[,2:ncol(X_tmp)]
         Y = Y_tmp[,2:ncol(Y_tmp)]

         print("Computing metrics for current epoch...")

         # Predict on the test dataset
         out = predict(X_test, Hin, Win, model, emas)
         accuracy_scalar = eval(out, Y_test)

         # Append to the epoch-wise metrics
         loss_metric[epoch, 1] = loss_avg
         accuracy[epoch, 1] = accuracy_scalar

         print("Epoch Avg. Loss: " + loss_avg)
         print("Epoch Accuracy: " + accuracy_scalar)
     }

     learned_model = model
     learned_emas = emas
 }

 predict = function(matrix[double] X, int Hin, int Win,
     list[unknown] model, list[unknown] emas)
     return(matrix[double] out) {
     /*
     * Computes the class probability predictions of a convolutional
     * net using the "ResNet" architecture.
     *
     * The input matrix, X, has N examples, each represented as a 3D
     * volume unrolled into a single vector.
     *
     * Inputs:
     *  - X: Input data matrix, of shape (N, C*Hin*Win).
     *
     * Outputs:
     *  - probs: Class probabilities, of shape (N, K).
     */

     # Predict on test dataset
     mode = "train"
     [out, temp_emas, temp_cached_out, temp_cached_means_vars] = resnet::forward(X, Hin, Win, model, mode, emas)
 }


 eval = function(matrix[double] probs, matrix[double] Y)
     return(double accuracy) {
     /*
     * Evaluates a convolutional net using the "ResNet" architecture.
     *
     * The probs matrix contains the class probability predictions
     * of K classes over N examples.  The targets, Y, have K classes,
     * and are one-hot encoded.
     *
     * Inputs:
     *  - probs: Class probabilities, of shape (N, K).
     *  - Y: Target matrix, of shape (N, K).
     *
     * Outputs:
     *  - accuracy: Scalar accuracy, of shape (1).
     */
     correct_pred = rowIndexMax(probs) == rowIndexMax(Y)
     accuracy = mean(correct_pred)
 }
	#-------------------------------------------------------------
	#
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.
	#
	#-------------------------------------------------------------

	# MNIST Resnet - Train
	#
	# This script trains a convolutional net using the "ResNet" architecture
	# on images of handwritten digits.
	#
	# Inputs:
	# - train: File containing labeled MNIST training images.
	# The format is "label, pixel_1, pixel_2, ..., pixel_n".
	# - test: File containing labeled MNIST test images.
	# The format is "label, pixel_1, pixel_2, ..., pixel_n".
	# - C: Number of color chanels in the images.
	# - Hin: Input image height.
	# - Win: Input image width.
	# - epochs: [DEFAULT: 10] Total number of full training loops over
	# the full data set.
	# - out_dir: [DEFAULT: "."] Directory to store weights and bias
	# matrices of trained model, as well as final test accuracy.
	# - fmt: [DEFAULT: "csv"] File format of `train` and `test` data.
	# Options include: "csv", "mm", "text", and "binary".
	#
	# Outputs:
	# - W1, W2, W3, W4: Files containing the trained weights of the model.
	# - b1, b2, b3, b4: Files containing the trained biases of the model.
	# - accuracy: File containing the accuracy and loss on the test data over all epochs.
	#
	# Data:
	# The MNIST dataset contains labeled images of handwritten digits,
	# where each example is a 28x28 pixel image of grayscale values in
	# the range [0,255] stretched out as 784 pixels, and each label is
	# one of 10 possible digits in [0,9].
	#
	# Sample Invocation (running from outside the `nn` folder):
	# 1. Download data (60,000 training examples, and 10,000 test examples)
	# ```
	# nn/examples/get_mnist_data.sh
	# ```
	#
	# 2. Execute using Spark
	# ```
	# spark-submit --master local[*] --driver-memory 10G
	# --conf spark.driver.maxResultSize=0 --conf spark.rpc.message.maxSize=128
	# $SYSTEMDS_ROOT/target/SystemDS.jar -f nn/examples/mnist_resnet.dml
	# -nvargs train=nn/examples/data/mnist/mnist_train.csv test=nn/examples/data/mnist/mnist_test.csv
	# C=1 Hin=28 Win=28 epochs=10 out_dir=nn/examples/model/mnist_resnet
	# ```
	#

	source("nn/networks/resnet18.dml") as resnet
	source("scripts/nn/layers/softmax_cross_entropy_loss.dml") as loss_nn

	# Read the data
	fmt = "csv"
	train = read("scripts/nn/examples/data/mnist_train.csv", format=fmt)
	test = read("scripts/nn/examples/data/mnist_test.csv", format=fmt)
	out_dir = "scripts/nn/example/model/mnist_resnet"

	# Extract images and labels
	images = train[,2:ncol(train)]
	labels = train[,1]
	images_test = test[,2:ncol(test)]
	labels_test = test[,1]
	classes = 10

	# Scale images to [-1,1], and one-hot encode the labels
	N = nrow(images)
	N_test = nrow(images_test)
	X = (images / 255.0) * 2 - 1
	X = cbind(X, X, X) # Resnet assumes C=3 so we duplicate the data along the channels
	Y = table(seq(1, N), labels+1, N, 10)
	X_test = (images_test / 255.0) * 2 - 1
	X_test = cbind(X_test, X_test, X_test)
	Y_test = table(seq(1, N_test), labels_test+1, N_test, 10)

	# Split into training (55,000 examples) and validation (5,000 examples)
	#X = images[5001:nrow(images),]
	#X_val = images[1:5000,]
	#Y = labels[5001:nrow(images),]
	#Y_val = labels[1:5000,]

	# Get initial model parameters
	[model, ema_means_vars] = resnet::init(classes, -1)

	# Get initial optimizer parameters
	optimizer_params = resnet::init_lars_optim_params(classes)
	# optimizer_params = resnet::init_sgd_momentum_optim_params(classes)
	# optimizer_params = resnet::init_adam_optim_params(classes)

	# Define image properties
	Hin = 28
	Win = 28
	#N_val = 0

	# Define training parameters
	epochs = 90
	batch_size = 512

	[accuracy, loss_metric, learned_model, learned_emas] = train(X, Y, X_test, Y_test, model, ema_means_vars, N, Hin, Win, epochs, batch_size, optimizer_params)

	write(accuracy, "scripts/nn/examples/out/resnet_mnist_accuracy.csv", format="csv")
	write(loss_metric, "scripts/nn/examples/out/resnet_mnist_loss.csv", format="csv")

	#Train
	train = function(matrix[double] X, matrix[double] Y, matrix[double] X_test, matrix[double] Y_test, list[unknown] model, list[unknown] emas, int samples, int Hin,
	int Win, int epochs, int batch_size, list[unknown] optim_params)
	return (matrix[double] accuracy, matrix[double] loss_metric,
	list[unknown] learned_model, list[unknown] learned_emas) {

	# --- LEARNING RATE SCHEDULE HYPERPARAMETERS ---
	# The learning rate we want to reach AFTER warmup
	initial_lr = 0.01
	# A very small final learning rate to decay towards
	end_lr = 0.0001
	# Number of warmup epochs, as per the paper
	warmup_epochs = 5
	# The exponent for the polynomial decay, as per the paper
	power = 2.0

	# Optimizer hyperparameters
	momentum = 0.9
	trust_coeff = 0.001
	weight_decay = 0.0001

	# Adam optimizer hyperparameters
	beta1 = 0.9
	beta2 = 0.999
	epsilon = 1e-8

	# Calculate total iterations for the schedule
	iterations_per_epoch = ceil(samples / batch_size)
	total_iterations = epochs * iterations_per_epoch
	warmup_iterations = warmup_epochs * iterations_per_epoch
	decay_iterations = total_iterations - warmup_iterations

	# Initialize metrics
	learned_model = list()
	learned_emas = list()
	accuracy = matrix(0, rows=epochs, cols=1)
	loss_metric = matrix(0, rows=epochs, cols=1)

	iterations = ceil(samples/batch_size)
	mode = "train"

	for (epoch in 1:epochs) {
	loss_avg = 0.0

	print("Start epoch: " + epoch + "/" + epochs)

	for (i in 1:iterations) {
	print(" - Iteration: " + i + "/" + iterations)

	# --- START DYNAMIC LEARNING RATE LOGIC ---
	current_iteration = (epoch - 1) * iterations_per_epoch + i
	current_lr = 0.0

	if (current_iteration < warmup_iterations) {
	# 1. Linear Warmup Phase
	# Linearly increase LR from 0 to initial_lr over warmup_iterations
	current_lr = initial_lr * (as.double(current_iteration) / warmup_iterations)
	} else {
	# 2. Polynomial Decay Phase
	decay_step = current_iteration - warmup_iterations
	decay_progress = as.double(decay_step) / decay_iterations
	current_lr = end_lr + (initial_lr - end_lr) * (1 - decay_progress)^power
	}

	if (i == 1) { # Print LR once per epoch to reduce log spam
	print("Using Learning Rate: " + current_lr)
	}
	# --- END DYNAMIC LEARNING RATE LOGIC ---

	# Get batch
	start = (i - 1) * batch_size + 1
	end = min(samples, i * batch_size)
	X_batch = X[start:end,]
	Y_batch = Y[start:end,]

	# Forward pass
	[out, emas, cached_out, cached_means_vars] = resnet::forward(X_batch, Hin, Win, model, mode, emas)

	# Loss
	loss = loss_nn::forward(out, Y_batch)
	if (i %% 10 == 0) { # Print loss less frequently on large datasets
	print(" - Iteration: " + i + "/" + iterations + ", Loss: " + loss)
	}
	loss_avg = (loss_avg * (i - 1) + loss) / i

	# Backward
	dOut = loss_nn::backward(out, Y_batch)
	[dX, gradients] = resnet::backward(dOut, cached_out, model, cached_means_vars)

	# Update parameters
	[model, optim_params] = resnet::update_params_with_lars(model, gradients, current_lr, momentum, weight_decay, trust_coeff,
	optim_params)
	# [model, optim_params] = resnet::update_params_with_sgd_momentum(model, gradients, current_lr, momentum, optim_params)

	# [model, optim_params] = resnet::update_params_with_adam(model, gradients, current_lr, beta1, beta2, epsilon, current_iteration, optim_params)
	}

	# Reshuffle mini batches
	r = rand(rows=nrow(Y), cols=1, min=0, max=1, pdf="uniform")
	X_tmp = order(target=cbind(r, X), by=1)
	Y_tmp = order(target=cbind(r, Y), by=1)
	X = X_tmp[,2:ncol(X_tmp)]
	Y = Y_tmp[,2:ncol(Y_tmp)]

	print("Computing metrics for current epoch...")

	# Predict on the test dataset
	out = predict(X_test, Hin, Win, model, emas)
	accuracy_scalar = eval(out, Y_test)

	# Append to the epoch-wise metrics
	loss_metric[epoch, 1] = loss_avg
	accuracy[epoch, 1] = accuracy_scalar

	print("Epoch Avg. Loss: " + loss_avg)
	print("Epoch Accuracy: " + accuracy_scalar)
	}

	learned_model = model
	learned_emas = emas
	}

	predict = function(matrix[double] X, int Hin, int Win,
	list[unknown] model, list[unknown] emas)
	return(matrix[double] out) {
	/*
	* Computes the class probability predictions of a convolutional
	* net using the "ResNet" architecture.
	*
	* The input matrix, X, has N examples, each represented as a 3D
	* volume unrolled into a single vector.
	*
	* Inputs:
	* - X: Input data matrix, of shape (N, CHinWin).
	*
	* Outputs:
	* - probs: Class probabilities, of shape (N, K).
	*/

	# Predict on test dataset
	mode = "train"
	[out, temp_emas, temp_cached_out, temp_cached_means_vars] = resnet::forward(X, Hin, Win, model, mode, emas)
	}


	eval = function(matrix[double] probs, matrix[double] Y)
	return(double accuracy) {
	/*
	* Evaluates a convolutional net using the "ResNet" architecture.
	*
	* The probs matrix contains the class probability predictions
	* of K classes over N examples. The targets, Y, have K classes,
	* and are one-hot encoded.
	*
	* Inputs:
	* - probs: Class probabilities, of shape (N, K).
	* - Y: Target matrix, of shape (N, K).
	*
	* Outputs:
	* - accuracy: Scalar accuracy, of shape (1).
	*/
	correct_pred = rowIndexMax(probs) == rowIndexMax(Y)
	accuracy = mean(correct_pred)
	}