From eb1af0b50e4e368daf5c4c5f9089a575adefcd49 Mon Sep 17 00:00:00 2001 From: cyberguli Date: Mon, 19 Feb 2024 23:09:10 +0100 Subject: [PATCH] add models and layers backward test --- tests/test_layers/test_conv.py | 33 ++++++++++++++ tests/test_layers/test_fourier.py | 36 +++++++++++++++ tests/test_layers/test_residual.py | 35 ++++++++++++++- tests/test_layers/test_spectral_conv.py | 36 +++++++++++++++ tests/test_model/test_deeponet.py | 32 +++++++++++++ tests/test_model/test_fnn.py | 9 ++++ tests/test_model/test_fno.py | 60 +++++++++++++++++++++++++ tests/test_model/test_mionet.py | 45 +++++++++++++++++++ tests/test_model/test_network.py | 12 +++++ tests/test_model/test_residualfnn.py | 11 +++++ 10 files changed, 308 insertions(+), 1 deletion(-) diff --git a/tests/test_layers/test_conv.py b/tests/test_layers/test_conv.py index f8ef437..8f322ac 100644 --- a/tests/test_layers/test_conv.py +++ b/tests/test_layers/test_conv.py @@ -106,6 +106,39 @@ def test_forward(): conv(x) +def test_backward(): + model = MLP + + x = torch.rand(dim_input) + x = make_grid(x) + x.requires_grad = True + # simple backward + conv = ContinuousConvBlock(channel_input, + channel_output, + dim, + stride, + model=model) + conv(x) + l=torch.mean(conv(x)) + l.backward() + assert x._grad.shape == torch.Size([2, 2, 20, 3]) + x = torch.rand(dim_input) + x = make_grid(x) + x.requires_grad = True + + # simple backward with optimization + conv = ContinuousConvBlock(channel_input, + channel_output, + dim, + stride, + model=model, + optimize=True) + conv(x) + l=torch.mean(conv(x)) + l.backward() + assert x._grad.shape == torch.Size([2, 2, 20, 3]) + + def test_transpose(): model = MLP diff --git a/tests/test_layers/test_fourier.py b/tests/test_layers/test_fourier.py index 826c445..f9c874b 100644 --- a/tests/test_layers/test_fourier.py +++ b/tests/test_layers/test_fourier.py @@ -20,6 +20,18 @@ def test_forward_1d(): sconv(x) +def test_backward_1d(): + sconv = FourierBlock1D(input_numb_fields=input_numb_fields, + output_numb_fields=output_numb_fields, + n_modes=4) + x = torch.rand(batch, input_numb_fields, 10) + x.requires_grad = True + sconv(x) + l = torch.mean(sconv(x)) + l.backward() + assert x._grad.shape == torch.Size([5, 3, 10]) + + def test_constructor_2d(): FourierBlock2D(input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, @@ -34,6 +46,18 @@ def test_forward_2d(): sconv(x) +def test_backward_2d(): + sconv = FourierBlock2D(input_numb_fields=input_numb_fields, + output_numb_fields=output_numb_fields, + n_modes=[5, 4]) + x = torch.rand(batch, input_numb_fields, 10, 10) + x.requires_grad = True + sconv(x) + l = torch.mean(sconv(x)) + l.backward() + assert x._grad.shape == torch.Size([5, 3, 10, 10]) + + def test_constructor_3d(): FourierBlock3D(input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, @@ -46,3 +70,15 @@ def test_forward_3d(): n_modes=[5, 4, 4]) x = torch.rand(batch, input_numb_fields, 10, 10, 10) sconv(x) + + +def test_backward_3d(): + sconv = FourierBlock3D(input_numb_fields=input_numb_fields, + output_numb_fields=output_numb_fields, + n_modes=[5, 4, 4]) + x = torch.rand(batch, input_numb_fields, 10, 10, 10) + x.requires_grad = True + sconv(x) + l = torch.mean(sconv(x)) + l.backward() + assert x._grad.shape == torch.Size([5, 3, 10, 10, 10]) diff --git a/tests/test_layers/test_residual.py b/tests/test_layers/test_residual.py index fc61b7f..03425a5 100644 --- a/tests/test_layers/test_residual.py +++ b/tests/test_layers/test_residual.py @@ -22,6 +22,17 @@ def test_forward_residual_block(): assert y.shape[1] == 3 assert y.shape[0] == x.shape[0] +def test_backward_residual_block(): + + res_block = ResidualBlock(input_dim=10, output_dim=3, hidden_dim=4) + + x = torch.rand(size=(80, 10)) + x.requires_grad = True + y = res_block(x) + l = torch.mean(y) + l.backward() + assert x._grad.shape == torch.Size([80,10]) + def test_constructor_no_activation_no_dropout(): linear_layer = nn.Linear(10, 20) enhanced_linear = EnhancedLinear(linear_layer) @@ -59,6 +70,17 @@ def test_forward_enhanced_linear_no_dropout(): assert y.shape[1] == 3 assert y.shape[0] == x.shape[0] +def test_backward_enhanced_linear_no_dropout(): + + enhanced_linear = EnhancedLinear(nn.Linear(10, 3)) + + x = torch.rand(size=(80, 10)) + x.requires_grad = True + y = enhanced_linear(x) + l = torch.mean(y) + l.backward() + assert x._grad.shape == torch.Size([80, 10]) + def test_forward_enhanced_linear_dropout(): enhanced_linear = EnhancedLinear(nn.Linear(10, 3), dropout=0.5) @@ -66,4 +88,15 @@ def test_forward_enhanced_linear_dropout(): x = torch.rand(size=(80, 10)) y = enhanced_linear(x) assert y.shape[1] == 3 - assert y.shape[0] == x.shape[0] \ No newline at end of file + assert y.shape[0] == x.shape[0] + +def test_backward_enhanced_linear_dropout(): + + enhanced_linear = EnhancedLinear(nn.Linear(10, 3), dropout=0.5) + + x = torch.rand(size=(80, 10)) + x.requires_grad = True + y = enhanced_linear(x) + l = torch.mean(y) + l.backward() + assert x._grad.shape == torch.Size([80, 10]) diff --git a/tests/test_layers/test_spectral_conv.py b/tests/test_layers/test_spectral_conv.py index 129d38b..3ff1ee3 100644 --- a/tests/test_layers/test_spectral_conv.py +++ b/tests/test_layers/test_spectral_conv.py @@ -20,6 +20,18 @@ def test_forward_1d(): sconv(x) +def test_backward_1d(): + sconv = SpectralConvBlock1D(input_numb_fields=input_numb_fields, + output_numb_fields=output_numb_fields, + n_modes=4) + x = torch.rand(batch, input_numb_fields, 10) + x.requires_grad = True + sconv(x) + l=torch.mean(sconv(x)) + l.backward() + assert x._grad.shape == torch.Size([5,3,10]) + + def test_constructor_2d(): SpectralConvBlock2D(input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, @@ -34,6 +46,18 @@ def test_forward_2d(): sconv(x) +def test_backward_2d(): + sconv = SpectralConvBlock2D(input_numb_fields=input_numb_fields, + output_numb_fields=output_numb_fields, + n_modes=[5, 4]) + x = torch.rand(batch, input_numb_fields, 10, 10) + x.requires_grad = True + sconv(x) + l=torch.mean(sconv(x)) + l.backward() + assert x._grad.shape == torch.Size([5,3,10,10]) + + def test_constructor_3d(): SpectralConvBlock3D(input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, @@ -46,3 +70,15 @@ def test_forward_3d(): n_modes=[5, 4, 4]) x = torch.rand(batch, input_numb_fields, 10, 10, 10) sconv(x) + + +def test_backward_3d(): + sconv = SpectralConvBlock3D(input_numb_fields=input_numb_fields, + output_numb_fields=output_numb_fields, + n_modes=[5, 4, 4]) + x = torch.rand(batch, input_numb_fields, 10, 10, 10) + x.requires_grad = True + sconv(x) + l=torch.mean(sconv(x)) + l.backward() + assert x._grad.shape == torch.Size([5,3,10,10,10]) diff --git a/tests/test_model/test_deeponet.py b/tests/test_model/test_deeponet.py index cfba614..78819e5 100644 --- a/tests/test_model/test_deeponet.py +++ b/tests/test_model/test_deeponet.py @@ -56,6 +56,21 @@ def test_forward_extract_int(): aggregator='*') model(data) +def test_backward_extract_int(): + data = torch.rand((20, 3)) + branch_net = FeedForward(input_dimensions=1, output_dimensions=10) + trunk_net = FeedForward(input_dimensions=2, output_dimensions=10) + model = DeepONet(branch_net=branch_net, + trunk_net=trunk_net, + input_indeces_branch_net=[0], + input_indeces_trunk_net=[1, 2], + reduction='+', + aggregator='*') + data.requires_grad = True + model(data) + l=torch.mean(model(data)) + l.backward() + assert data._grad.shape == torch.Size([20,3]) def test_forward_extract_str_wrong(): branch_net = FeedForward(input_dimensions=1, output_dimensions=10) @@ -68,3 +83,20 @@ def test_forward_extract_str_wrong(): aggregator='*') with pytest.raises(RuntimeError): model(data) + +def test_backward_extract_str_wrong(): + data = torch.rand((20, 3)) + branch_net = FeedForward(input_dimensions=1, output_dimensions=10) + trunk_net = FeedForward(input_dimensions=2, output_dimensions=10) + model = DeepONet(branch_net=branch_net, + trunk_net=trunk_net, + input_indeces_branch_net=['a'], + input_indeces_trunk_net=['b', 'c'], + reduction='+', + aggregator='*') + data.requires_grad = True + with pytest.raises(RuntimeError): + model(data) + l=torch.mean(model(data)) + l.backward() + assert data._grad.shape == torch.Size([20,3]) diff --git a/tests/test_model/test_fnn.py b/tests/test_model/test_fnn.py index bdd38fe..d02dcb8 100644 --- a/tests/test_model/test_fnn.py +++ b/tests/test_model/test_fnn.py @@ -35,3 +35,12 @@ def test_forward(): fnn = FeedForward(dim_in, dim_out) output_ = fnn(data) assert output_.shape == (data.shape[0], dim_out) + +def test_backward(): + dim_in, dim_out = 3, 2 + fnn = FeedForward(dim_in, dim_out) + data.requires_grad = True + output_ = fnn(data) + l=torch.mean(output_) + l.backward() + assert data._grad.shape == torch.Size([20,3]) diff --git a/tests/test_model/test_fno.py b/tests/test_model/test_fno.py index 322ce1f..3c8094b 100644 --- a/tests/test_model/test_fno.py +++ b/tests/test_model/test_fno.py @@ -60,6 +60,24 @@ def test_1d_forward(): assert out.shape == torch.Size([batch_size, resolution[0], output_channels]) +def test_1d_backward(): + input_channels = 1 + input_ = torch.rand(batch_size, resolution[0], input_channels) + lifting_net = torch.nn.Linear(input_channels, lifting_dim) + projecting_net = torch.nn.Linear(60, output_channels) + fno = FNO(lifting_net=lifting_net, + projecting_net=projecting_net, + n_modes=5, + dimensions=1, + inner_size=60, + n_layers=2) + input_.requires_grad = True + out = fno(input_) + l = torch.mean(out) + l.backward() + assert input_.grad.shape == torch.Size([batch_size, resolution[0], input_channels]) + + def test_2d_forward(): input_channels = 2 input_ = torch.rand(batch_size, resolution[0], resolution[1], @@ -77,6 +95,27 @@ def test_2d_forward(): [batch_size, resolution[0], resolution[1], output_channels]) +def test_2d_backward(): + input_channels = 2 + input_ = torch.rand(batch_size, resolution[0], resolution[1], + input_channels) + lifting_net = torch.nn.Linear(input_channels, lifting_dim) + projecting_net = torch.nn.Linear(60, output_channels) + fno = FNO(lifting_net=lifting_net, + projecting_net=projecting_net, + n_modes=5, + dimensions=2, + inner_size=60, + n_layers=2) + input_.requires_grad = True + out = fno(input_) + l = torch.mean(out) + l.backward() + assert input_.grad.shape == torch.Size([ + batch_size, resolution[0], resolution[1], input_channels + ]) + + def test_3d_forward(): input_channels = 3 input_ = torch.rand(batch_size, resolution[0], resolution[1], resolution[2], @@ -93,3 +132,24 @@ def test_3d_forward(): assert out.shape == torch.Size([ batch_size, resolution[0], resolution[1], resolution[2], output_channels ]) + + +def test_3d_backward(): + input_channels = 3 + input_ = torch.rand(batch_size, resolution[0], resolution[1], resolution[2], + input_channels) + lifting_net = torch.nn.Linear(input_channels, lifting_dim) + projecting_net = torch.nn.Linear(60, output_channels) + fno = FNO(lifting_net=lifting_net, + projecting_net=projecting_net, + n_modes=5, + dimensions=3, + inner_size=60, + n_layers=2) + input_.requires_grad = True + out = fno(input_) + l = torch.mean(out) + l.backward() + assert input_.grad.shape == torch.Size([ + batch_size, resolution[0], resolution[1], resolution[2], input_channels + ]) diff --git a/tests/test_model/test_mionet.py b/tests/test_model/test_mionet.py index 4e9c03c..174251e 100644 --- a/tests/test_model/test_mionet.py +++ b/tests/test_model/test_mionet.py @@ -36,6 +36,22 @@ def test_forward_extract_str(): model(input_) +def test_backward_extract_str(): + data = torch.rand((20, 3)) + data.requires_grad = True + input_vars = ['a', 'b', 'c'] + input_ = LabelTensor(data, input_vars) + branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10) + branch_net2 = FeedForward(input_dimensions=1, output_dimensions=10) + trunk_net = FeedForward(input_dimensions=1, output_dimensions=10) + networks = {branch_net1: ['a'], branch_net2: ['b'], trunk_net: ['c']} + model = MIONet(networks=networks, reduction='+', aggregator='*') + model(input_) + l = torch.mean(model(input_)) + l.backward() + assert data._grad.shape == torch.Size([20,3]) + + def test_forward_extract_int(): branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10) branch_net2 = FeedForward(input_dimensions=1, output_dimensions=10) @@ -45,6 +61,20 @@ def test_forward_extract_int(): model(data) +def test_backward_extract_int(): + data = torch.rand((20, 3)) + data.requires_grad = True + branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10) + branch_net2 = FeedForward(input_dimensions=1, output_dimensions=10) + trunk_net = FeedForward(input_dimensions=1, output_dimensions=10) + networks = {branch_net1: [0], branch_net2: [1], trunk_net: [2]} + model = MIONet(networks=networks, reduction='+', aggregator='*') + model(data) + l = torch.mean(model(data)) + l.backward() + assert data._grad.shape == torch.Size([20,3]) + + def test_forward_extract_str_wrong(): branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10) branch_net2 = FeedForward(input_dimensions=1, output_dimensions=10) @@ -53,3 +83,18 @@ def test_forward_extract_str_wrong(): model = MIONet(networks=networks, reduction='+', aggregator='*') with pytest.raises(RuntimeError): model(data) + + +def test_backward_extract_str_wrong(): + data = torch.rand((20, 3)) + data.requires_grad = True + branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10) + branch_net2 = FeedForward(input_dimensions=1, output_dimensions=10) + trunk_net = FeedForward(input_dimensions=1, output_dimensions=10) + networks = {branch_net1: ['a'], branch_net2: ['b'], trunk_net: ['c']} + model = MIONet(networks=networks, reduction='+', aggregator='*') + with pytest.raises(RuntimeError): + model(data) + l = torch.mean(model(data)) + l.backward() + assert data._grad.shape == torch.Size([20,3]) diff --git a/tests/test_model/test_network.py b/tests/test_model/test_network.py index 20e514d..870480e 100644 --- a/tests/test_model/test_network.py +++ b/tests/test_model/test_network.py @@ -35,3 +35,15 @@ def test_forward(): with pytest.raises(AssertionError): net(data) + +def test_backward(): + net = Network(model=torchmodel, + input_variables=['x', 'y', 'z'], + output_variables=['a', 'b', 'c', 'd'], + extra_features=None) + data = torch.rand((20, 3)) + data.requires_grad = True + out = net.torchmodel(data) + l = torch.mean(out) + l.backward() + assert data._grad.shape == torch.Size([20, 3]) \ No newline at end of file diff --git a/tests/test_model/test_residualfnn.py b/tests/test_model/test_residualfnn.py index eef05b1..1c0cbf8 100644 --- a/tests/test_model/test_residualfnn.py +++ b/tests/test_model/test_residualfnn.py @@ -24,3 +24,14 @@ def test_forward(): x = torch.rand(10, 2) model = ResidualFeedForward(input_dimensions=2, output_dimensions=1) model(x) + + +def test_backward(): + x = torch.rand(10, 2) + x.requires_grad = True + model = ResidualFeedForward(input_dimensions=2, output_dimensions=1) + model(x) + l = torch.mean(model(x)) + l.backward() + assert x.grad.shape == torch.Size([10, 2]) + \ No newline at end of file