Movatterモバイル変換


[0]ホーム

URL:


Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Sign up
Appearance settings

Zuko density estimators#1088

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
Show file tree
Hide file tree
Changes from1 commit
Commits
Show all changes
27 commits
Select commitHold shift + click to select a range
aad9335
update zuko to 1.1.0
bkmiMar 18, 2024
c879f7a
test zuko_gmm commit
Mar 18, 2024
447b39c
build_zuko_nsf added
Mar 19, 2024
1ee8221
add build_zuko_naf, update test
bkmiMar 19, 2024
f42bee5
add license change to pr template.
janfbMar 18, 2024
be43e0b
CLN pyproject.toml (#1009)
tomMoralMar 19, 2024
caa87a9
fix x_o and broken link tutorial 7 (#1003)
MatthijspalsMar 19, 2024
7f2d205
replace prepare_for_sbi in tutorials (#1013)
zinaStefMar 19, 2024
a88a0f0
add zuko density estimators
Mar 20, 2024
c6c82ae
modify file with ruff
Mar 20, 2024
bd4a013
not working gmm
Mar 21, 2024
8af200d
update tests for PR
Mar 22, 2024
daef025
update PR for pyright
Mar 22, 2024
4bc42f4
resolve pyright
Mar 22, 2024
b300eae
add reportArgumentType
Mar 22, 2024
11e15b1
resolve pyright issue
Mar 22, 2024
fb1394b
resolve all issues pyright
Mar 22, 2024
f6afdc4
resolve pyright
Mar 22, 2024
f928df3
add typing and docstring
Mar 27, 2024
c58157d
add functions from factory to test
Mar 27, 2024
27a8af3
remove comment mdn file
Mar 27, 2024
2559e0c
add docstrings flow file
Mar 27, 2024
ccd63ba
add docstring in density_estimator_test.py
Mar 27, 2024
a553ac4
Update sbi/neural_nets/flow.py
anastasiakrouglovaMar 27, 2024
e4bd842
Update sbi/neural_nets/flow.py
anastasiakrouglovaMar 27, 2024
c4dbf6f
Update sbi/neural_nets/flow.py
anastasiakrouglovaMar 27, 2024
5882301
removed pyright
Mar 27, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
PrevPrevious commit
NextNext commit
build_zuko_nsf added
  • Loading branch information
Nastya Krouglova authored andNastya Krouglova committedMar 19, 2024
commit447b39c20c44f1ecf6cf842a10af1af293be1843
126 changes: 120 additions & 6 deletionssbi/neural_nets/flow.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -429,27 +429,141 @@ def mask_in_layer(i):

#################################### Functions Nastya #############################################

# Doesn't work yet
def build_zuko_gmm(
batch_x: Tensor,
batch_y: Tensor,
z_score_x: Optional[str] = "independent",
z_score_y: Optional[str] = "independent",
hidden_features: Union[Sequence[int], int] = 50,
num_transforms: int = 5,
# hidden_features: Union[Sequence[int], int] = 50,
num_components: int = 10,
# num_transforms: int = 5,
embedding_net: nn.Module = nn.Identity(),
residual: bool = True,
randperm: bool = False,
#residual: bool = True,
#randperm: bool = False,
**kwargs,
) -> ZukoFlow:

pass
x_numel = batch_x[0].numel()
# Infer the output dimensionality of the embedding_net by making a forward pass.
check_data_device(batch_x, batch_y)
check_embedding_net_device(embedding_net=embedding_net, datum=batch_y)
y_numel = embedding_net(batch_y[:1]).numel()
if x_numel == 1:
warn(
"In one-dimensional output space, this flow is limited to Gaussians",
stacklevel=1,
)

# if isinstance(hidden_features, int):
# hidden_features = [hidden_features] * num_transforms

if x_numel == 1:
gmm = zuko.flows.GMM(
features=x_numel,
context=y_numel,
# hidden_features=hidden_features,
components=num_components,
)

# components = gmm.transform.transforms
# z_score_x_bool, structured_x = z_score_parser(z_score_x)
# if z_score_x_bool:
# # transforms = transforms
# transforms = (
# *transforms,
# standardizing_transform(batch_x, structured_x, backend="zuko"),
# )

z_score_y_bool, structured_y = z_score_parser(z_score_y)
if z_score_y_bool:
# Prepend standardizing transform to y-embedding.
embedding_net = nn.Sequential(
standardizing_net(batch_y, structured_y), embedding_net
)


neural_net = zuko.flows.Flow(num_components, gmm.shapes)
flow = ZukoFlow(neural_net, embedding_net, condition_shape=batch_y[0].shape)

return flow



def build_zuko_nsf(
batch_x: Tensor,
batch_y: Tensor,
z_score_x: Optional[str] = "independent",
z_score_y: Optional[str] = "independent",
hidden_features: int = 50,
num_transforms: int = 5,
num_bins: int = 10,
embedding_net: nn.Module = nn.Identity(),
tail_bound: float = 3.0,
hidden_layers_spline_context: int = 1,
num_blocks: int = 2,
dropout_probability: float = 0.0,
use_batch_norm: bool = False,
residual: bool = True, # params for batch learning
randperm: bool = False, # params for batch learning
**kwargs,
) -> ZukoFlow:

x_numel = batch_x[0].numel()
# Infer the output dimensionality of the embedding_net by making a forward pass.
check_data_device(batch_x, batch_y)
check_embedding_net_device(embedding_net=embedding_net, datum=batch_y)
y_numel = embedding_net(batch_y[:1]).numel()


if isinstance(hidden_features, int):
hidden_features = [hidden_features] * num_transforms

if x_numel == 1:
nsf = zuko.flows.NSF(
features=x_numel,
context=y_numel,
hidden_features=hidden_features,
transforms=num_transforms,
)
else:
nsf = zuko.flows.NSF(
features=x_numel,
context=y_numel,
hidden_features=hidden_features,
transforms=num_transforms,
randperm=randperm,
residual=residual,
)


transforms = nsf.transform.transforms


z_score_x_bool, structured_x = z_score_parser(z_score_x)
if z_score_x_bool:
# transforms = transforms
transforms = (
*transforms,
standardizing_transform(batch_x, structured_x, backend="zuko"),
)

z_score_y_bool, structured_y = z_score_parser(z_score_y)
if z_score_y_bool:
# Prepend standardizing transform to y-embedding.
embedding_net = nn.Sequential(
standardizing_net(batch_y, structured_y), embedding_net
)

neural_net = zuko.flows.Flow(transforms, nsf.base)
flow = ZukoFlow(neural_net, embedding_net, condition_shape=batch_y[0].shape)


return flow



####################################################################################################
# ####################################################################################################


def build_zuko_maf(
Expand Down
23 changes: 20 additions & 3 deletionstests/density_estimator_test.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,8 +9,7 @@
from torch.distributions import MultivariateNormal

from sbi.neural_nets.density_estimators import NFlowsFlow, ZukoFlow
from sbi.neural_nets.flow import build_nsf, build_zuko_maf

from sbi.neural_nets.flow import build_nsf, build_zuko_nsf, build_zuko_maf, build_zuko_gmm

@pytest.mark.parametrize("density_estimator", (NFlowsFlow, ZukoFlow))
@pytest.mark.parametrize("input_dims", (1, 2))
Expand DownExpand Up@@ -53,13 +52,29 @@ def forward(self, x):
embedding_net=EmbeddingNet(),
)
elif density_estimator == ZukoFlow:
estimator = build_zuko_maf(
# estimator = build_zuko_gmm(
# batch_input,
# batch_context,
# hidden_features=10,
# num_components=2,
# embedding_net=EmbeddingNet(),
# )
estimator = build_zuko_nsf(
batch_input,
batch_context,
hidden_features=10,
num_transforms=2,
embedding_net=EmbeddingNet(),
)
# estimator = build_zuko_maf(
# batch_input,
# batch_context,
# hidden_features=10,
# num_transforms=2,
# embedding_net=EmbeddingNet(),
# )



# Loss is only required to work for batched inputs and contexts
loss = estimator.loss(batch_input, batch_context)
Expand DownExpand Up@@ -213,3 +228,5 @@ def forward(self, x):
nsamples_test,
), f"log_prob shape is not correct. It is of shape {log_probs.shape}, but should \
be {(1, batch_context.shape[0], 2, nsamples_test)}"



[8]ページ先頭

©2009-2025 Movatter.jp