add mis coef net

client.py

@@ -88,18 +88,20 @@ def send_radiance(self):
         lum = np.stack((y,) * 3, axis=-1)
         scale = np.array([3, 3, 3])
         lum = lum / scale
+        pdf_brdf = np.random.random((self.num_points, 1)).astype(
+            np.float32)
         self.points_data = np.concatenate(
-            (self.points_data, lum.reshape([len(lum), 3])), axis=1, dtype=np.float32
+            (self.points_data, lum.reshape([len(lum), 3]), torch.from_numpy(pdf_brdf)), axis=1, dtype=np.float32
         )
         self.client_socket.send(
             len(
-                self.points_data[:, [14, 15, 16, 0, 1, 2, 3, 4, 5, 6, 7]].tobytes()
+                self.points_data[:, [14, 15, 16, 0, 1, 2, 3, 4, 5, 6, 7, 17]].tobytes()
             ).to_bytes(4, "little")
         )  # bytes
 
         raw_data = bytearray()
         raw_data.extend(
-            self.points_data[:, [14, 15, 16, 0, 1, 2, 3, 4, 5, 6, 7]].tobytes()
+            self.points_data[:, [14, 15, 16, 0, 1, 2, 3, 4, 5, 6, 7, 17]].tobytes()
         )  # send r,g,b x, y, z, norm_x, norm_y, norm_z, dir_x, dir_y, dir_z
         self.client_socket.send(raw_data)
 

experiment_config.py

@@ -31,8 +31,8 @@ class ExperimentConfig:
     ndims: int = 3
     funcname: str = "Gaussian"
     coupling_name: str = "piecewiseQuadratic"
-    num_context_features: int = 0
-    hidden_dim: int = 10
+    num_context_features: int = 8
+    hidden_dim: int = 256
     n_hidden_layers: int = 3
     network_type: str = "MLP"
     blob: Union[int, None] = None
@@ -48,6 +48,7 @@ class ExperimentConfig:
     max_train_buffer_size: int = 2_000_000
     features_mode: str = "all_features"  # 'no_features' 'xyz' 'all_features'
     one_bounce_mode: bool = True
+    drop_zero: bool = False
 
     save_plots: blob = True
     plot_dimension: int = 2
@@ -89,6 +90,7 @@ def init_from_pyhocon(cls, pyhocon_config: pyhocon_wrapper.ConfigTree):
                 "train.features_mode", "all_features"
             ),
             one_bounce_mode=pyhocon_config.get_bool("train.one_bounce_mode", True),
+            drop_zero=pyhocon_config.get_bool("train.drop_zero", False),
             funcname=pyhocon_config.get_string("train.function"),
             coupling_name=pyhocon_config.get_string("train.coupling_name"),
             num_context_features=pyhocon_config.get_int("train.num_context_features"),

integrator.py

@@ -19,6 +19,7 @@ def __init__(
         self,
         func,
         flow,
+        coef_net,
         dist,
         optimizer,
         scheduler=None,
@@ -42,6 +43,7 @@ def __init__(
         self.global_step = 0
         self.scheduler_step = 0
         self.flow = flow.to(device)
+        self.coef_net = coef_net.to(device)
         self.dist: torch.distributions.Distribution = dist
         self.optimizer = optimizer
         self.scheduler = scheduler
@@ -227,10 +229,13 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False):
         self.flow.eval()
         if self.features_mode == "all_features":
             flow_context = torch.tensor(context[:, :8]).to(self.device)
+            flow_context = torch.Tensor(context[:, [0, 1, 2, 6, 7]]).to(self.device)
         elif self.features_mode == "xyz":
             flow_context = torch.tensor(context[:, :3]).to(self.device)
-        else:
+        elif self.features_mode == "no_features":
             flow_context = None
+        else:
+            raise AttributeError(f"expected features_mode [all_features, xyz, no_features], got {self.features_mode}")
         if inverse:
             z = torch.tensor(context[:, 8:]).to(
                 self.device
@@ -239,7 +244,7 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False):
             z[:, 1] = np.abs(np.cos(z[:, 1].cpu()))  # theta
             with torch.no_grad():
                 x, absdet = self.flow.inverse(
-                    z, context=torch.tensor(flow_context).to(self.device)
+                    z, context=flow_context
                 )
             return 1 / absdet.to("cpu")
         else:
@@ -252,7 +257,8 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False):
             )
             with torch.no_grad():
                 x, absdet = self.flow(z, context=flow_context)
-            return (x.to("cpu"), absdet.to("cpu"))
+                coef = self.coef_net(torch.tensor(context[:, [0, 1, 2, 6, 7]]).to(self.device), context=None)  # xyz , w0
+            return (x.to("cpu"), absdet.to("cpu"), coef.to("cpu"))
 
     def train_with_context(
         self,
@@ -273,73 +279,96 @@ def train_with_context(
 
         # Process #
         if self.features_mode == "all_features":
-            context_x = torch.Tensor(context[:, :-1]).to(self.device)
+            context_x = torch.Tensor(context[:, :-2]).to(self.device)
+            context_x = torch.Tensor(context[:, [0, 1, 2, 6, 7]]).to(self.device)
         elif self.features_mode == "xyz":
             context_x = torch.Tensor(context[:, :3]).to(self.device)
         else:
             context_x = None
-        context_y = context[:, -1]
+        context_y = context[:, 9]
         train_result = []
         start = time.time()
         logging.info(f"context_x time: {time.time() - start}")
         start = time.time()
 
-        context_y = torch.Tensor(context_y)
+        context_y = torch.Tensor(context_y).to(self.device)
+        context_с = torch.Tensor(context[:, [0, 1, 2, 6, 7]]).to(self.device)
+        hybrid_pdf = torch.Tensor(context[:, 8]).to(self.device)
         logging.info(f"context_y time: {time.time() - start}")
         start = time.time()
-        for batch_x, batch_y, batch_z in [(context_x, context_y, z)]:
-            logging.info(f"batch time: {time.time() - start}")
+        logging.info(f"batch time: {time.time() - start}")
 
-            start = time.time()
-            x, absdet = self.flow(batch_z, batch_x)
+        start = time.time()
+        x, absdet = self.flow(z, context_x)
+        coef = self.coef_net(context_с, context=None).reshape(absdet.size())  # xyz , w0
+        absdet.requires_grad_(True)
+        # y = self._func(x)
+        logging.info(f"flow time: {time.time() - start}")
+        # absdet *= torch.exp(log_prob) # P_X(x) = PZ(f^-1(x)) |det(df/dx)|^-1
 
-            absdet.requires_grad_(True)
-            # y = self._func(x)
-            logging.info(f"flow time: {time.time() - start}")
-            # absdet *= torch.exp(log_prob) # P_X(x) = PZ(f^-1(x)) |det(df/dx)|^-1
+        # --------------- START TODO compute loss ---------------
+        start = time.time()
+        y = context_y.to(self.device)
+        y = y + 1e-10
+        cross = torch.nn.CrossEntropyLoss()
+        loss = cross(absdet, y)
+        kl_loss = torch.nn.KLDivLoss()
+        kl_loss = self.loss_func
+        # loss = kl_loss(torch.log_softmax(absdet, -1), torch.softmax(y, -1))
+        # loss = kl_loss(absdet, y)
+        # loss = (
+        #         (y/absdet) * (torch.log(y) - torch.log(absdet))
+        # ).mean()
+        loss = (y*torch.log(y)).mean() - (y * torch.log(absdet)).mean()
+        loss.backward()
 
-            # --------------- START TODO compute loss ---------------
-            start = time.time()
-            y = batch_y.to(self.device)
+        absdet_coef = coef * absdet + (1 - coef) * hybrid_pdf
+        # loss_coef = kl_loss(torch.log_softmax(absdet_coef, -1), torch.softmax(y, -1))
+        loss_coef = (y*torch.log(y)).mean() - (y * torch.log(absdet_coef)).mean()
 
-            cross = torch.nn.CrossEntropyLoss()
-            loss = cross(absdet, y)
-            mean = torch.mean(y / absdet)
-            var = torch.var(y / absdet)
 
-            if var == 0:
-                var += np.finfo(np.float32).eps
+        tao = 0.5
+        b_tao = 0.5*((1/3)**(5*tao))
+        final_loss = b_tao*loss + (1-b_tao)*loss_coef
+        # final_loss.backward()
+
+        mean = torch.mean(y / absdet)
+        var = torch.var(y / absdet)
+
+        if var == 0:
+            var += np.finfo(np.float32).eps
+
+        # Backprop #
+        # loss = self.loss_func(y, absdet)
+        logging.info(f"loss time: {time.time() - start}")
 
-            # Backprop #
-            # loss = self.loss_func(y, absdet)
-            logging.info(f"loss time: {time.time() - start}")
+        start = time.time()
+        # loss.backward()
+        logging.info(f"backward time: {time.time() - start}")
+        print("\t" "Loss = %0.8f" % loss)
+        print("\t" "mean coef = %0.8f" % coef.mean())
+        # --------------- END TODO compute loss ---------------
 
+        if apply_optimizer:
             start = time.time()
-            loss.backward()
-            logging.info(f"backward time: {time.time() - start}")
-            print("\t" "Loss = %0.8f" % loss)
-            # --------------- END TODO compute loss ---------------
-
-            if apply_optimizer:
-                start = time.time()
-
-                self.apply_optimizer()
-                logging.info(f"optimizer time: {time.time() - start}")
-
-            # Integral #
-            return_dict = {"loss": loss.to("cpu").item(), "epoch": self.global_step}
-            self.global_step += 1
-            if lr:
-                return_dict["lr"] = self.optimizer.param_groups[0]["lr"]
-            if integral:
-                return_dict["mean"] = mean.to("cpu").item()
-                return_dict["uncertainty"] = (
-                    torch.sqrt(var / (context.shape[0] - 1.0)).to("cpu").item()
-                )
-            if points:
-                return_dict["z"] = batch_z.to("cpu")
-                return_dict["x"] = x.to("cpu")
-            train_result.append(return_dict)
+
+            self.apply_optimizer()
+            logging.info(f"optimizer time: {time.time() - start}")
+
+        # Integral #
+        return_dict = {"loss": loss.to("cpu").item(), "epoch": self.global_step}
+        self.global_step += 1
+        if lr:
+            return_dict["lr"] = self.optimizer.param_groups[0]["lr"]
+        if integral:
+            return_dict["mean"] = mean.to("cpu").item()
+            return_dict["uncertainty"] = (
+                torch.sqrt(var / (context.shape[0] - 1.0)).to("cpu").item()
+            )
+        if points:
+            return_dict["z"] = z.to("cpu")
+            return_dict["x"] = x.to("cpu")
+        train_result.append(return_dict)
         return train_result
 
     def apply_optimizer(self):

nis.py

@@ -31,6 +31,8 @@ def __init__(self, _config: ExperimentConfig):
         mode: ['server', 'experiment'] - if mode==server, don't use dimension reduction and function in visualization
         """
         self.config = _config
+        if self.config.gradient_accumulation:
+            self.config.num_training_steps = 1
         self.logs_dir = os.path.join("logs", self.config.experiment_dir_name)
         self.integrator = None
         self.visualize_object = None
@@ -55,7 +57,7 @@ def initialize(self, mode="server"):
         # masks = self.create_binary_mask(self.config.ndims)
         masks = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1]]
         if self.config.features_mode == "all_features":
-            num_context_features = self.config.num_context_features
+            num_context_features = self.config.num_context_features - 3
         elif self.config.features_mode == "xyz":
             num_context_features = 3
         else:
@@ -75,6 +77,14 @@ def initialize(self, mode="server"):
                 for mask in masks
             ]
         )
+        coef_net = UNet(
+                in_features=5,
+                out_features=1,
+                max_hidden_features=self.config.hidden_dim,
+                num_layers=self.config.n_hidden_layers,
+                nonlinearity=nn.ReLU(),
+                output_activation=nn.Sigmoid(),
+            )
         dist = torch.distributions.uniform.Uniform(
             torch.tensor([0.0] * self.config.ndims),
             torch.tensor([1.0] * self.config.ndims),
@@ -88,6 +98,7 @@ def initialize(self, mode="server"):
             scheduler=None,
             loss_func=self.config.loss_func,
             features_mode=self.config.features_mode,
+            coef_net=coef_net,
         )
 
         self.means = []
@@ -223,12 +234,16 @@ def binary_list(inval, length):
         return masks
 
     def get_samples(self, context):
-        # pdf_light_sample = self.integrator.sample_with_context(context, inverse=True)
-        [samples, pdf] = self.integrator.sample_with_context(context, inverse=False)
-        pdf_light_sample = torch.ones(pdf.size())
-        return [samples, pdf_light_sample, pdf]
+        pdf_light_sample = self.integrator.sample_with_context(context, inverse=True)
+        [samples, pdf, coef] = self.integrator.sample_with_context(context, inverse=False)
+        # pdf_light_sample = torch.ones(pdf.size())
+        return [samples, pdf_light_sample, pdf, coef]
 
     def train(self, context):
+        if self.config.drop_zero:
+            context = context[np.nonzero(context[:, -1])]
+        if context.size == 0:
+            return None
         z = torch.stack(self.integrator.generate_z_by_context(context))
         context = torch.tensor(context)
         if self.z_buffer is None:
@@ -247,6 +262,8 @@ def train(self, context):
                 -self.config.max_train_buffer_size:
             ]
             self.z_buffer = self.z_buffer[-self.config.max_train_buffer_size:]
+        if self.config.gradient_accumulation and self.train_sampling_call_difference != 1:
+            return
         for epoch in range(self.config.num_training_steps):
             start = time.time()
             indices = torch.randperm(len(self.context_buffer))[
@@ -261,7 +278,6 @@ def train(self, context):
                 integral=True,
                 points=True,
                 batch_size=self.config.batch_size,
-                apply_optimizer=not self.config.gradient_accumulation,
             )
 
             logging.info(f"Epoch time: {time.time() - start}")
@@ -275,7 +291,8 @@ def train(self, context):
                         self.log_tensorboard_train_step(epoch_result)
 
         if self.config.gradient_accumulation:
-            self.integrator.apply_optimizer()
+            self.z_buffer = None
+            self.context_buffer = None
 
         if self.train_sampling_call_difference == 1:
             self.integrator.z_mapper = {}  # Be careful with z_mapper