Commit e092a2aa authored by Johan Gudmundsson's avatar Johan Gudmundsson
Browse files

torch_proxies test pickle to save

parent e8970248

src/borch/module.py

+3 −2
Original line number Diff line number Diff line
@@ -292,7 +292,6 @@ class Observed(_Module): 
        )







class Module(_Module):

    """Acts as a ``torch.nn.Module`` but handles ``borch.RandomVariable`` s correctly.
@@ -453,7 +452,9 @@ class Module(_Module): 
        for k, v in tuple(self.observed.items()):

            if v is None:

                del self.observed[k]

                getattr(self.posterior, k)()  # redraw this sample

                rv = getattr(self.posterior, k, None)

                if rv is not None:

                    rv()  # redraw this sample



    def __getattr__(self, name):

        try:

src/borch/nn/borchify.py

+17 −7
Original line number Diff line number Diff line
@@ -26,8 +26,9 @@ from borch.utils.namespace_tools import ( 
    create_module,

)



_PICKLE_CACHE=create_module('PICKLE_CACHE', 'Cache of pickled modules')

_PICKLE_CACHE_PATH = f"{__name__}._PICKLE_CACHE"

PICKLE_CACHE = create_module("PICKLE_CACHE", "Cache of pickled modules")

_PICKLE_CACHE_PATH = f"{__name__}.PICKLE_CACHE"





def _get_priors_from_kwargs_(kwargs):

    priors = {}
@@ -58,6 +59,9 @@ def borch_proxy_class( 
    """

    Create a Bayesian version of a `torch.nn.Module`



    Note that if the modules exists in `torch.nn` they will be replaced with the

    equivalent module from `borch.nn` instead of creating a new class.



    Args:

        cls: an uninstanciated class that is a subclass of `torch.nn.Module`

        get_rv_factory: function that takes a string as an argument and returns
@@ -93,7 +97,7 @@ def borch_proxy_class( 
    new_cls = type(

        cls.__name__, (cls, BorchModule, *extra_inheriors), {"__init__": _init}

    )

    assign_docs(new_cls, cls, doc_prefix, False)

    assign_docs(new_cls, cls, doc_prefix)

    caller = caller if caller is not None else _PICKLE_CACHE_PATH

    new_cls.__module__ = caller

    return new_cls
@@ -112,6 +116,9 @@ def borch_classes( 
    """

    Create a Bayesian version of classes in a python module.



    Note that if the modules exists in `torch.nn` they will be replaced with the

    equivalent module from `borch.nn` instead of creating a new class.



    Args:

        module: a python module that contains `torch.nn.Module`s you want a

                Bayesian version of.
@@ -153,6 +160,9 @@ def borchify_namespace( 
    """

    Create a new module that contains bayesian versions of the `torch.nn.Module`s.



    Note that if the modules exists in `torch.nn` they will be replaced with the

    equivalent module from `borch.nn` instead of creating a new class.



    Args:

        module: a python module that contains `torch.nn.Module`s you want a

                Bayesian version of.
@@ -174,7 +184,7 @@ def borchify_namespace( 
        >>> bnn = borchify_namespace(torch.nn)

        >>> blinear = bnn.Linear(1,2)

        >>> type(blinear)

        <class 'borch.nn.borchify.Linear'>

        <class 'borch.nn.torch_proxies.Linear'>

    """

    mappings = borch_classes(

        module=module,
@@ -254,7 +264,7 @@ def borchify_module( 
        >>> linear = torch.nn.Linear(3, 3)  # create a linear module

        >>> blinear = borchify_module(linear)

        >>> type(blinear)

        <class 'borch.nn.borchify.Linear'>

        <class 'borch.nn.torch_proxies.Linear'>

    """

    # pylint: disable=unexpected-keyword-arg
@@ -327,9 +337,9 @@ def borchify_network( 
      >>> net = Net()

      >>> bnet = borchify_network(net)

      >>> type(bnet)

      <class 'borch.nn.borchify.Net'>

      <class 'borch.nn.borchify.PICKLE_CACHE.Net'>

      >>> type(bnet.linear)

      <class 'borch.nn.borchify.Linear'>

      <class 'borch.nn.torch_proxies.Linear'>

    """

    if cache is None:

        cache = {}

tests/borch/nn/test_torch_proxies.py

+50 −18
Original line number Diff line number Diff line
@@ -264,7 +264,6 @@ def test_sequential(): 
    assert loss > 0







def test_transformer_samples():

    transformer_model = borch.nn.Transformer(nhead=16, num_encoder_layers=12)

    src = torch.rand((10, 32, 512))
@@ -273,14 +272,14 @@ def test_transformer_samples(): 
    borch.sample(transformer_model)

    assert out != transformer_model(src, tgt).sum()





def test_transformer_submodules_are_borch_modules():

    transformer_model = borch.nn.Transformer(nhead=2, num_encoder_layers=2)

    for name, mod in transformer_model.named_children():

    for mod in transformer_model.children():

        if mod not in transformer_model.internal_modules:

            assert isinstance(mod, borch.Module)







@pytest.mark.parametrize(

    "rnn,state",

    [
@@ -302,69 +301,101 @@ def test_rnn_sum(rnn, state): 
        assert [x.grad.abs().sum() > 0 is not None for x in rnn.parameters()]







@pytest.fixture(params = [

@pytest.fixture(

    params=[

        (borch.nn.Linear(2, 3), torch.randn(3, 2)),

        (borch.nn.Conv1d(6, 33, 3, stride=2), torch.randn(2, 6, 10)),

        (borch.nn.Conv2d(6, 33, 3, stride=2), torch.randn(2, 6, 10, 10)),

        (borch.nn.Conv3d(2, 33, 3, stride=2), torch.randn(1, 2, 10, 10, 3)),

        (borch.nn.ConvTranspose2d(6, 33, (3, 5), stride=(2, 1), padding=(4, 2)), torch.randn(2, 6, 10, 10)),

        (

            borch.nn.ConvTranspose2d(6, 33, (3, 5), stride=(2, 1), padding=(4, 2)),

            torch.randn(2, 6, 10, 10),

        ),

        (borch.nn.Bilinear(20, 30, 40), (torch.randn(128, 20), torch.randn(128, 30))),

        (borch.nn.GRU(10, 20, 2), (torch.randn(2, 3, 10), torch.randn(2, 3, 20))),

        (borch.nn.RNN(10, 20, 2), (torch.randn(2, 3, 10), torch.randn(2, 3, 20))),

        (borch.nn.LSTM(10, 20, 2), (torch.randn(2, 3, 10), (torch.randn(2, 3, 20), torch.randn(2, 3, 20)))),

        (

            borch.nn.LSTM(10, 20, 2),

            (torch.randn(2, 3, 10), (torch.randn(2, 3, 20), torch.randn(2, 3, 20))),

        ),

        (borch.nn.RNNCell(10, 20), (torch.randn(3, 10), torch.randn(3, 20))),

        (borch.nn.GRUCell(10, 20), (torch.randn(3, 10), torch.randn(3, 20))),

        (borch.nn.LSTMCell(10, 20), (torch.randn(3, 10), (torch.randn(3, 20), torch.randn(3, 20)))),

        (borch.nn.EmbeddingBag(10, 3, mode='sum'), (torch.tensor([1,2,4,5,4,3,2,9], dtype=torch.long), torch.tensor([0,4], dtype=torch.long))),

        (

            borch.nn.LSTMCell(10, 20),

            (torch.randn(3, 10), (torch.randn(3, 20), torch.randn(3, 20))),

        ),

        (

            borch.nn.EmbeddingBag(10, 3, mode="sum"),

            (

                torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long),

                torch.tensor([0, 4], dtype=torch.long),

            ),

        ),

        (borch.nn.Embedding(3, 5, max_norm=True), torch.tensor([1, 2])),

    ])

    ]

)

def mod_ipt(request):

    return request.param



def _torch_save_module(mod_ipt):



def _save_module(save, load, mod_ipt):

    mod, ipt = mod_ipt

    buffer = BytesIO()

    torch.save(mod, buffer)

    save(mod, buffer)

    buffer.seek(0)

    new= torch.load(buffer)

    new = load(buffer)

    assert isinstance(new, borch.Module)

    assert isinstance(new, type(mod))

    return new, ipt



def _torch_save_state_dict(mod_ipt):



def _save_state_dict(save, load, mod_ipt):

    mod, ipt = mod_ipt

    buffer = BytesIO()

    torch.save(mod.state_dict(), buffer)

    save(mod.state_dict(), buffer)

    buffer.seek(0)

    sd= torch.load(buffer)

    sd = load(buffer)

    borch.sample(mod)

    _run_sum(mod, ipt)

    mod.load_state_dict(sd)

    return mod, ipt





_torch_save_module = lambda x: _save_module(torch.save, torch.load, x)

_pkl_save_module = lambda x: _save_module(pkl.dump, pkl.load, x)

_torch_save_state_dict = lambda x: _save_state_dict(torch.save, torch.load, x)

_pkl_save_state_dict = lambda x: _save_state_dict(pkl.dump, pkl.load, x)





def _sample(mod_ipt):

    mod, ipt = mod_ipt

    borch.sample(mod)

    return mod, ipt





@pytest.fixture(params = [

@pytest.fixture(

    params=[

        lambda x: x,

        _torch_save_module,

        _torch_save_state_dict,

        _sample,

        lambda x: _sample(_torch_save_module(x)),

        _pkl_save_module,

        _pkl_save_state_dict,

        lambda x: _sample(_pkl_save_module(x)),

        lambda x: _sample(_pkl_save_state_dict(x)),

        lambda x: _torch_save_module(_sample(x)),

        lambda x: _sample(_torch_save_module(_sample(x))),

        lambda x: _sample(_torch_save_state_dict(x)),

        lambda x: _sample(_torch_save_state_dict(_sample(x))),

        lambda x: _torch_save_state_dict(_sample(x)),

    ])

    ]

)

def mod_inpt_transform(mod_ipt, request):

    fn = request.param

    return fn(mod_ipt)





def _run_sum(mod, ipt):

    if isinstance(ipt, tuple):

        out = mod(*ipt)
@@ -374,6 +405,7 @@ def _run_sum(mod, ipt): 
        out = out[0]

    return out.sum()





def test_samples_propperly(mod_inpt_transform):

    mod, ipt = mod_inpt_transform

    out = _run_sum(mod, ipt)

tests/borch/test_module.py

+5 −0
Original line number Diff line number Diff line
@@ -114,6 +114,11 @@ def test_unobserve(model): 
    assert model.normal != val





def test_unobserve_with_missing_observed_rvs(model):

    model.observe(missing=torch.randn(2))

    module.unobserve(model)





def test_observe_is_part_of_state_dict(model):

    val = test_observe(model)

    assert model.state_dict()["observed.normal"] == val