use LabelTensor, fix minor, docs
This commit is contained in:
58
tests/test_fnn.py
Normal file
58
tests/test_fnn.py
Normal file
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
import pytest
|
||||
|
||||
from pina import LabelTensor
|
||||
from pina.model import FeedForward
|
||||
|
||||
class myFeature(torch.nn.Module):
|
||||
"""
|
||||
Feature: sin(pi*x)
|
||||
"""
|
||||
def __init__(self):
|
||||
super(myFeature, self).__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return torch.sin(torch.pi * x.extract('a'))
|
||||
|
||||
|
||||
data = torch.rand((20, 3))
|
||||
input_vars = ['a', 'b', 'c']
|
||||
output_vars = ['d', 'e']
|
||||
input_ = LabelTensor(data, input_vars)
|
||||
|
||||
|
||||
def test_constructor():
|
||||
FeedForward(input_vars, output_vars)
|
||||
FeedForward(3, 4)
|
||||
FeedForward(input_vars, output_vars, extra_features=[myFeature()])
|
||||
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():
|
||||
fnn = FeedForward(input_vars, output_vars)
|
||||
output_ = fnn(input_)
|
||||
assert output_.labels == output_vars
|
||||
|
||||
|
||||
def test_forward2():
|
||||
dim_in, dim_out = 3, 2
|
||||
fnn = FeedForward(dim_in, dim_out)
|
||||
output_ = fnn(input_)
|
||||
assert output_.shape == (input_.shape[0], dim_out)
|
||||
|
||||
|
||||
def test_forward_features():
|
||||
fnn = FeedForward(input_vars, output_vars, extra_features=[myFeature()])
|
||||
output_ = fnn(input_)
|
||||
assert output_.labels == output_vars
|
||||
Reference in New Issue
Block a user