Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix to PeriodicPadding to match pytorch circular padding #2435

Merged
merged 4 commits into from
Apr 18, 2024
Merged

Fix to PeriodicPadding to match pytorch circular padding #2435

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
95806c9
Fix to PeriodicPadding to match pytorch circular padding
jvwilliams23 Apr 15, 2024
6b0387b
Minor changes to docstring and kwargs
jvwilliams23 Apr 15, 2024
6f06e37
Fix test by comparing against pytorch F.pad with mode="circular"
jvwilliams23 Apr 16, 2024
e04fefe
Skip test if pytorch isn't available
fiedorowicz1 Apr 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
303 changes: 46 additions & 257 deletions ci_test/unit_tests/test_unit_module_periodic_padding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
30 changes: 15 additions & 15 deletions python/lbann/modules/transformations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand All @@ -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],
Expand All @@ -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)
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 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
Expand All @@ -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
Expand All @@ -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],
Expand All @@ -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
Expand Down