add models and layers backward test

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
 

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])

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,11 +70,33 @@ 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)
 
     x = torch.rand(size=(80, 10))
     y = enhanced_linear(x)
     assert y.shape[1] == 3
-    assert y.shape[0] == x.shape[0]
+    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])

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])

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])

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])

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
+    ])

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])