From 95806c911a27e3f63d1115dbe3d66432d149a3db Mon Sep 17 00:00:00 2001 From: jvwilliams23 Date: Mon, 15 Apr 2024 16:50:00 +0100 Subject: [PATCH 1/4] Fix to PeriodicPadding to match pytorch circular padding --- python/lbann/modules/transformations.py | 28 ++++++++++++------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/python/lbann/modules/transformations.py b/python/lbann/modules/transformations.py index 5241ebd5edc..4f2fbdaeab0 100644 --- a/python/lbann/modules/transformations.py +++ b/python/lbann/modules/transformations.py @@ -90,15 +90,15 @@ def Cumsum(x, dims, axis=0): def PeriodicPadding2D(x, height, width, padding=1): - """ For 2D images of the shape (B, *, height, width) + """ For 2D images of the shape (channels, height, width) Args: - x (lbann.Layer): input tensor to padded of the shape (*, height, width) - height (int): 2nd dimension of the 4D tensor - width (int): 3rd dimension of the 4D tensor - padding (int): The amount to pad (default: 1) + x (lbann.Layer): input tensor to padded of the shape (channels, height, width) + height (int): 1st dimension of the 3D tensor + width (int): 2nd dimension of the 3D tensor + padding (int): The amount to pad on each side (default: 1) returns: (lbann.Layer): Returns periodically padded layer of - shape (*, height + padding, width + padding) + shape (channels, height + 2 * padding, width + 2 * padding) """ horizontal_slices = lbann.Slice(x, slice_points=[0, padding, height - padding, height], @@ -107,7 +107,7 @@ def PeriodicPadding2D(x, height, width, padding=1): _ = lbann.Identity(horizontal_slices) bottom = lbann.Identity(horizontal_slices) - x = lbann.Concatenation([top, x, bottom], axis=1) + x = lbann.Concatenation([bottom, x, top], axis=1) vertical_slices = lbann.Slice(x, slice_points=[0, padding, width - padding, width], @@ -116,21 +116,21 @@ def PeriodicPadding2D(x, height, width, padding=1): _ = lbann.Identity(vertical_slices) right = lbann.Identity(vertical_slices) - x = lbann.Concatenation([left, x, right], axis=2) + x = lbann.Concatenation([right, x, left], axis=2) return x def PeriodicPadding3D(x, depth, height, width, padding=1, name=None): - """ For 3D volumes of the shape (B, *, channel, depth, height, width) + """ For 3D volumes of the shape (B, channels, depth, height, width) Args: - x (lbann.Layer): input tensor to padded of the shape (*, channel, depth, height, width) + x (lbann.Layer): input tensor to be padded of the shape (channels, depth, height, width) depth (int): 1st dimension of the 4D tensor height (int): 2nd dimension of the 4D tensor width (int): 3rd dimension of the 4D tensor padding (int): The amount to pad (default: 1) returns: (lbann.Layer): Returns periodically padded layer of - shape (*, depth + padding, height + padding, width + padding) + shape (channels, depth + 2 * padding, height + 2 * padding, width + 2 * padding) """ # To do: Hack to get around slice and concatenation limitation. Remove when # support for arbitrary dimensional slice + concatenation is added @@ -142,7 +142,7 @@ def PeriodicPadding3D(x, depth, height, width, padding=1, name=None): _ = lbann.Identity(depth_slices) d2 = lbann.Identity(depth_slices) - x = lbann.Concatenation([d1, x, d2], axis=1) + x = lbann.Concatenation([d2, x, d1], axis=1) # To do: Hack to get around slice and concatenation limitation. Remove when # support for arbitrary dimensional slice + concatenation is added @@ -154,7 +154,7 @@ def PeriodicPadding3D(x, depth, height, width, padding=1, name=None): _ = lbann.Identity(height_slices) h2 = lbann.Identity(height_slices) - x = lbann.Concatenation([h1, x, h2], axis=1) + x = lbann.Concatenation([h2, x, h1], axis=1) width_slices = lbann.Slice(x, slice_points=[0, padding, width - padding, width], @@ -163,7 +163,7 @@ def PeriodicPadding3D(x, depth, height, width, padding=1, name=None): _ = lbann.Identity(width_slices) w2 = lbann.Identity(width_slices) - x = lbann.Concatenation([w1, x, w2], axis=2) + x = lbann.Concatenation([w2, x, w1], axis=2) # To do: Hack to get around slice and concatenation limitation. Remove when # support for arbitrary dimensional slice + concatenation is added From 6b0387bac32f3107150026febc89c9983197ac97 Mon Sep 17 00:00:00 2001 From: jvwilliams23 Date: Mon, 15 Apr 2024 17:02:54 +0100 Subject: [PATCH 2/4] Minor changes to docstring and kwargs As batch size is implicit, remove from docstring. Remove unused "name" kwarg in function. --- python/lbann/modules/transformations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/lbann/modules/transformations.py b/python/lbann/modules/transformations.py index 4f2fbdaeab0..6f85b74d659 100644 --- a/python/lbann/modules/transformations.py +++ b/python/lbann/modules/transformations.py @@ -120,8 +120,8 @@ def PeriodicPadding2D(x, height, width, padding=1): return x -def PeriodicPadding3D(x, depth, height, width, padding=1, name=None): - """ For 3D volumes of the shape (B, channels, depth, height, width) +def PeriodicPadding3D(x, depth, height, width, padding=1): + """ For 3D volumes of the shape (channels, depth, height, width) Args: x (lbann.Layer): input tensor to be padded of the shape (channels, depth, height, width) depth (int): 1st dimension of the 4D tensor From 6f06e373f86c7401d2375ccf2d09d2ba26240aaa Mon Sep 17 00:00:00 2001 From: jvwilliams23 Date: Tue, 16 Apr 2024 08:56:00 +0100 Subject: [PATCH 3/4] Fix test by comparing against pytorch F.pad with mode="circular" --- .../test_unit_module_periodic_padding.py | 303 +++--------------- 1 file changed, 46 insertions(+), 257 deletions(-) diff --git a/ci_test/unit_tests/test_unit_module_periodic_padding.py b/ci_test/unit_tests/test_unit_module_periodic_padding.py index 300459489b8..354b5b78f27 100644 --- a/ci_test/unit_tests/test_unit_module_periodic_padding.py +++ b/ci_test/unit_tests/test_unit_module_periodic_padding.py @@ -1,258 +1,47 @@ -import functools +import lbann +import lbann.modules as lm import numpy as np -import os -import os.path -import sys -import operator -import math -from lbann.modules.transformations import PeriodicPadding3D, PeriodicPadding2D -# Bamboo utilities -current_file = os.path.realpath(__file__) -current_dir = os.path.dirname(current_file) -sys.path.insert(0, os.path.join(os.path.dirname(current_dir), 'common_python')) -import tools - -# ============================================== -# Objects for Python data reader -# ============================================== -# Note: The Python data reader imports this file as a module and calls -# the functions below to ingest data. - - -def make_random_array(shape, seed): - """Hacked function to generate a random array. - - NumPy's RNG produces different values with different NumPy - versions. This function is helpful when array values must be - identical across all runs, e.g. when checking against precomputed - metric values. - - Args: - shape (Iterable of int): Array dimensions - seed (int): Parameter for RNG. Must be non-zero. - Returns: - numpy.ndarray: Array of `np.float32`. Values will be in - [-0.5,0.5). - - """ - size = functools.reduce(operator.mul, shape) - eps = np.finfo(np.float32).eps - x = (seed / np.linspace(math.sqrt(eps), 0.1, size)) % 1 - 0.5 - return x.reshape(shape).astype(np.float32) - - -# Data -_num_samples = 23 -_sample_dims = [6, 11, 7] -_sample_dims_3d = [2, 3, 11, 7] -_sample_size = functools.reduce(operator.mul, _sample_dims) -_samples = make_random_array([_num_samples] + _sample_dims, 7) - - -# Sample access functions -def get_sample(index): - return _samples[index, :].reshape(-1) - - -def num_samples(): - return _num_samples - - -def sample_dims(): - return (_sample_size,) - -# ============================================== -# Periodic Padding -# ============================================== - - -def periodic_padding_2D(data, padding): - """ - Args: - data (np.array) : Input array of shape (B, c, h, w) - padding (int) : Amount of padding around data - Returns - (np.array): Padded atensor with shape - (B, c, h+2*padding, h+2*padding) - """ - _, c, h, w = data.shape - top_slice = data[:, :, :padding, :] - bottom_slice = data[:, :, h - padding:, :] - inter = np.concatenate((bottom_slice, data, top_slice), axis=2) - left_slice = inter[:, :, :, :padding] - right_slice = inter[:, :, :, w - padding:] - return np.concatenate((right_slice, inter, left_slice), axis=3) - - -def periodic_padding_3D(data, padding): - """ - Args: - data (np.array) : Input array of shape (B, c, d, h, w) - padding (int) : Amount of padding around data - Returns - (np.array): Padded atensor with shape - (B, c, d+2*padding, h+2*padding, h+2*padding) - """ - _, c, d, h, w = data.shape - d_slice_start = data[:, :, :padding, :, :] - d_slice_end = data[:, :, d - padding:, :, :] - inter = np.concatenate((d_slice_end, data, d_slice_start), axis=2) - h_slice_start = inter[:, :, :, :padding, :] - h_slice_end = inter[:, :, :, h - padding:, :] - - inter = np.concatenate((h_slice_end, inter, h_slice_start), axis=3) - - w_slice_start = inter[:, :, :, :, :padding] - w_slice_end = inter[:, :, :, :, w - padding:] - - return_val = np.concatenate((w_slice_end, inter, w_slice_start), axis=4) - return return_val - - -def setup_experiment(lbann, weekly): - """Construct LBANN experiment. - - Args: - lbann (module): Module for LBANN Python frontend - - """ - mini_batch_size = num_samples() // 2 - trainer = lbann.Trainer(mini_batch_size) - model = construct_model(lbann) - data_reader = construct_data_reader(lbann) - optimizer = lbann.NoOptimizer() - return trainer, model, data_reader, optimizer, None # Don't request any specific number of nodes - - -def construct_model(lbann): - """Construct LBANN model. - - Args: - lbann (module): Module for LBANN Python frontend - - """ - - # Input data - # Note: Sum with a weights layer so that gradient checking will - # verify that error signals are correct. - x_weights = lbann.Weights(optimizer=lbann.SGD(), - initializer=lbann.ConstantInitializer(value=0.0)) - x = lbann.Sum(lbann.Reshape(lbann.Input(data_field='samples'), - dims=_sample_dims), - lbann.WeightsLayer(weights=x_weights, - dims=_sample_dims)) - x_lbann = x - - # Objects for LBANN model - obj = [] - metrics = [] - callbacks = [] - - x_2D = lbann.Reshape(x_lbann, - dims=_sample_dims) - y = PeriodicPadding2D(x_2D, - _sample_dims[1], - _sample_dims[2], - padding=2) - z = lbann.L2Norm2(y) - obj.append(z) - metrics.append(lbann.Metric(z, name="Padding_2D")) - - x_np = _samples - y_np = periodic_padding_2D(x_np, padding=2) - z_np = tools.numpy_l2norm2(y_np) / _num_samples - tol = 8 * z_np * np.finfo(np.float32).eps - - metric_callback_2d = lbann.CallbackCheckMetric(metric=metrics[-1].name, - lower_bound=z_np - tol, - upper_bound=z_np + tol, - error_on_failure=True, - execution_modes='test') - - x_3D = lbann.Reshape(x_lbann, - dims=_sample_dims_3d) - y = PeriodicPadding3D(x_3D, - _sample_dims_3d[1], - _sample_dims_3d[2], - _sample_dims_3d[3], - padding=1) - z = lbann.L2Norm2(y) - obj.append(z) - metrics.append(lbann.Metric(z, name="Padding_3D")) - x_np = _samples.reshape([_num_samples] + _sample_dims_3d) - y_np = periodic_padding_3D(x_np, padding=1) - z_np = tools.numpy_l2norm2(y_np) / _num_samples - - tol = 8 * z_np * np.finfo(np.float32).eps - - metric_callback_3d = lbann.CallbackCheckMetric(metric=metrics[-1].name, - lower_bound=z_np - tol, - upper_bound=z_np + tol, - error_on_failure=True, - execution_modes='test') - metrics.append(lbann.Metric(z, name="Padding_3D")) - - # ------------------------------------------ - # Gradient checking - # ------------------------------------------ - - callbacks.append(lbann.CallbackCheckGradients(error_on_failure=True)) - - # ------------------------------------------ - # Construct model - # ------------------------------------------ - - num_epochs = 0 - return lbann.Model(num_epochs, - layers=lbann.traverse_layer_graph(x_lbann), - objective_function=obj, - metrics=metrics, - callbacks=callbacks) - - -def construct_data_reader(lbann): - """Construct Protobuf message for Python data reader. - - The Python data reader will import the current Python file to - access the sample access functions. - - Args: - lbann (module): Module for LBANN Python frontend - - """ - - # Note: The training data reader should be removed when - # https://github.com/LLNL/lbann/issues/1098 is resolved. - message = lbann.reader_pb2.DataReader() - message.reader.extend([ - tools.create_python_data_reader( - lbann, - current_file, - 'get_sample', - 'num_samples', - 'sample_dims', - 'train' - ) - ]) - message.reader.extend([ - tools.create_python_data_reader( - lbann, - current_file, - 'get_sample', - 'num_samples', - 'sample_dims', - 'test' - ) - ]) - return message - - -# ============================================== -# Setup PyTest -# ============================================== - -# Create test functions that can interact with PyTest -# Note: Create test name by removing ".py" from file name -_test_name = os.path.splitext(os.path.basename(current_file))[0] -for _test_func in tools.create_tests(setup_experiment, _test_name, skip_clusters=["catalyst"]): - globals()[_test_func.__name__] = _test_func +import test_util +import pytest +from torch import Tensor +import torch.nn.functional as F + +@test_util.lbann_test(check_gradients=False) +def test_periodic_padding_2D(): + # Prepare reference output + np.random.seed(20240228) + shape = [1, 4, 16, 20] + _, _, height, width = shape + x = np.random.rand(*shape).astype(np.float32) + ref = F.pad(Tensor(x), (1,1,1,1), mode="circular").numpy() + + tester = test_util.ModelTester() + + x = tester.inputs(x)[0] + reference = tester.make_reference(ref) + # Test layer + y = lm.PeriodicPadding2D(x, height=height, width=width, padding=1) + + # Set test loss + tester.set_loss(lbann.MeanSquaredError(y, reference)) + return tester + +@test_util.lbann_test(check_gradients=False) +def test_periodic_padding_3D(): + # Prepare reference output + np.random.seed(20240228) + shape = [1, 4, 8, 16, 20] + _, _, depth, height, width = shape + x = np.random.rand(*shape).astype(np.float32) + ref = F.pad(Tensor(x), (1,1,1,1,1,1), mode="circular").numpy() + + tester = test_util.ModelTester() + + x = tester.inputs(x)[0] + reference = tester.make_reference(ref) + # Test layer + y = lm.PeriodicPadding3D(x, depth=depth, height=height, width=width, padding=1) + + # Set test loss + tester.set_loss(lbann.MeanSquaredError(y, reference)) + return tester From e04fefec80b6e5b8da9928a7029beba4d2d9c5a0 Mon Sep 17 00:00:00 2001 From: Pier Fiedorowicz Date: Tue, 16 Apr 2024 18:07:34 -0700 Subject: [PATCH 4/4] Skip test if pytorch isn't available --- ci_test/unit_tests/test_unit_module_periodic_padding.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ci_test/unit_tests/test_unit_module_periodic_padding.py b/ci_test/unit_tests/test_unit_module_periodic_padding.py index 354b5b78f27..24f0b1c1fe1 100644 --- a/ci_test/unit_tests/test_unit_module_periodic_padding.py +++ b/ci_test/unit_tests/test_unit_module_periodic_padding.py @@ -6,6 +6,13 @@ from torch import Tensor import torch.nn.functional as F +try: + from torch import Tensor + import torch.nn.functional as F +except: + pytest.skip("PyTorch is required to run this test.", allow_module_level=True) + + @test_util.lbann_test(check_gradients=False) def test_periodic_padding_2D(): # Prepare reference output