51 lines
1.2 KiB
Python
51 lines
1.2 KiB
Python
import torch
|
|
import pytest
|
|
|
|
from pina.model import FeedForward
|
|
|
|
data = torch.rand((20, 3))
|
|
input_vars = 3
|
|
output_vars = 4
|
|
|
|
|
|
def test_constructor():
|
|
FeedForward(input_vars, output_vars)
|
|
FeedForward(input_vars, output_vars, inner_size=10, n_layers=20)
|
|
FeedForward(input_vars, output_vars, layers=[10, 20, 5, 2])
|
|
FeedForward(
|
|
input_vars, output_vars, layers=[10, 20, 5, 2], func=torch.nn.ReLU
|
|
)
|
|
FeedForward(
|
|
input_vars,
|
|
output_vars,
|
|
layers=[10, 20, 5, 2],
|
|
func=[torch.nn.ReLU, torch.nn.ReLU, None, torch.nn.Tanh],
|
|
)
|
|
|
|
|
|
def test_constructor_wrong():
|
|
with pytest.raises(RuntimeError):
|
|
FeedForward(
|
|
input_vars,
|
|
output_vars,
|
|
layers=[10, 20, 5, 2],
|
|
func=[torch.nn.ReLU, torch.nn.ReLU],
|
|
)
|
|
|
|
|
|
def test_forward():
|
|
dim_in, dim_out = 3, 2
|
|
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])
|