Upload flow_pgm.py

pgm/flow_pgm.py

@@ -28,6 +28,23 @@ class Hparams:
             setattr(self, k, v)
 
 
+def is_one_hot(tensor):
+    """
+    Check if the given tensor is a valid one-hot tensor.
+
+    Args:
+    tensor (torch.Tensor): A tensor to check.
+
+    Returns:
+    bool: True if tensor is one-hot, False otherwise.
+    """
+    if tensor.ndim != 2:
+        return False
+
+    # Check if there is exactly one '1' in each row and all other elements are '0'
+    return torch.all((tensor.sum(dim=1) == 1) & (tensor.max(dim=1).values == 1) & (tensor.min(dim=1).values == 0))
+
+
 class BasePGM(nn.Module):
     def __init__(self):
         super().__init__()
@@ -461,13 +478,13 @@ class ChestPGM(BasePGM):
         self.variables = {
             "race": "categorical",
             "sex": "binary",
-            "finding": "binary",
+            "finding": "categorical",
             "age": "continuous",
         }
         # Discrete variables that are not root nodes
-        self.discrete_variables = {"finding": "binary"}
+        self.discrete_variables = {"finding": "categorical"}
         # define base distributions
-        for k in ["a"]:  # , "f"]:
+        for k in ["a", "f"]:
             self.register_buffer(f"{k}_base_loc", torch.zeros(1))
             self.register_buffer(f"{k}_base_scale", torch.ones(1))
         # age spline flow
@@ -481,26 +498,51 @@ class ChestPGM(BasePGM):
                 # self.age_constraints,
             ]
         )
-        # Finding (conditional) via MLP, a -> f
-        finding_net = DenseNN(1, [8, 16], param_dims=[2], nonlinearity=nn.Sigmoid())
+        # Finding (conditional) via MLP, a -> fss
+        finding_net = DenseNN(1, [8, 16], param_dims=[3], nonlinearity=nn.Softmax())
         self.finding_transform_GumbelMax = ConditionalGumbelMax(
             context_nn=finding_net, event_dim=0
         )
         # log space for sex and race
         self.sex_logit = nn.Parameter(np.log(1 / 2) * torch.ones(1))
         self.race_logits = nn.Parameter(np.log(1 / 3) * torch.ones(1, 3))
+        self.finding_logits = nn.Parameter(np.log(1 / 3) * torch.ones(1, 3))
 
-        input_shape = (args.input_channels, args.input_res, args.input_res)
+        if args.setup != "sup_pgm":
+            from resnet import CustomBlock, ResNet, ResNet18
 
-        if args.enc_net == "cnn":
+            shared_model = ResNet(
+                CustomBlock,
+                layers=[2, 2, 2, 2],
+                widths=[64, 128, 256, 512],
+                norm_layer=lambda c: nn.GroupNorm(min(32, c // 4), c),
+            )
+            # shared_model = torchvision.models.resnet18(weights=None)
+            shared_model.conv1 = nn.Conv2d(
+                args.input_channels,
+                64,
+                kernel_size=7,
+                stride=2,
+                padding=3,
+                bias=False,
+            )
+            kwargs = {
+                "in_shape": (args.input_channels, *(args.input_res,) * 2),
+                "base_model": shared_model,
+            }
             # q(s | x) ~ Bernoulli(f(x))
-            self.encoder_s = CNN(input_shape, num_outputs=1)
-            # q(r | x) ~ OneHotCategorical(logits=f(x))
-            self.encoder_r = CNN(input_shape, num_outputs=3)
+            self.encoder_s = ResNet18(num_outputs=1, **kwargs)
+            # q(r | x) ~ OneHotCategorical(f(x))
+            self.encoder_r = ResNet18(num_outputs=3, **kwargs)
             # q(f | x) ~ Bernoulli(f(x))
-            self.encoder_f = CNN(input_shape, num_outputs=1)
+            self.encoder_f = ResNet18(num_outputs=3, **kwargs)
             # q(a | x, f) ~ Normal(mu(x), sigma(x))
-            self.encoder_a = CNN(input_shape, num_outputs=1, context_dim=1)
+            self.encoder_a = ResNet18(num_outputs=2, context_dim=3, **kwargs)
+            self.f = (
+                lambda x: args.std_fixed * torch.ones_like(x)
+                if args.std_fixed > 0
+                else F.softplus(x)
+            )
 
     def model(self) -> Dict[str, Tensor]:
         pyro.module("ChestPGM", self)
@@ -521,11 +563,14 @@ class ChestPGM(BasePGM):
 
         # p(f | a), finding as OneHotCategorical conditioned on age
         # finding_dist_base = dist.Gumbel(self.f_base_loc, self.f_base_scale).to_event(1)
-        finding_dist_base = dist.Gumbel(torch.zeros(1), torch.ones(1)).to_event(1)
-
+        # finding_dist = ConditionalTransformedDistributionGumbelMax(
+        #     finding_dist_base, [self.finding_transform_GumbelMax]
+        # ).condition(age)
+        finding_dist_base = dist.OneHotCategorical(logits=self.finding_logits) #.to_event(1)
+        # finding_dist_base = dist.RelaxedOneHotCategoricalStraightThrough(temperature =1e-5,logits=self.finding_logits) #.to_event(1)
         finding_dist = ConditionalTransformedDistributionGumbelMax(
-            finding_dist_base, [self.finding_transform_GumbelMax]
-        ).condition(age)
+            finding_dist_base, [self.finding_transform_GumbelMax]).condition(age)
+            
         finding = pyro.sample("finding", finding_dist)
 
         return {
@@ -548,8 +593,8 @@ class ChestPGM(BasePGM):
                 pyro.sample("race", qr_x)
             # q(f | x)
             if obs["finding"] is None:
-                f_prob = torch.sigmoid(self.encoder_f(obs["x"]))
-                qf_x = dist.Bernoulli(probs=f_prob).to_event(1)
+                f_prob = F.softmax(self.encoder_f(obs["x"]),dim=-1)
+                qf_x = dist.OneHotCategorical(probs=f_prob).to_event(1)
                 obs["finding"] = pyro.sample("finding", qf_x)
             # q(a | x, f)
             if obs["age"] is None:
@@ -576,8 +621,9 @@ class ChestPGM(BasePGM):
             pyro.sample("race_aux", qr_x, obs=obs["race"])
 
             # q(f | x)
-            f_prob = torch.sigmoid(self.encoder_f(obs["x"]))
-            qf_x = dist.Bernoulli(probs=f_prob).to_event(1)
+            f_probs = F.softmax(self.encoder_f(obs["x"]), dim=-1)
+            qf_x= dist.OneHotCategorical(probs=f_probs)  # .to_event(1)
+            # with pyro.poutine.scale(scale=0.5):
             pyro.sample("finding_aux", qf_x, obs=obs["finding"])
 
             # q(a | x, f)
@@ -594,7 +640,7 @@ class ChestPGM(BasePGM):
         # q(r | x)
         r_probs = F.softmax(self.encoder_r(obs["x"]), dim=-1)
         # q(f | x)
-        f_prob = torch.sigmoid(self.encoder_f(obs["x"]))
+        f_prob = F.softmax(self.encoder_f(obs["x"]),dim=-1)
         # q(a | x, f)
         a_loc, _ = self.encoder_a(obs["x"], y=obs["finding"]).chunk(2, dim=-1)