From 33dc911ff9e317d238a6182351fce059f4e16be9 Mon Sep 17 00:00:00 2001 From: alexpetersburg Date: Wed, 8 Feb 2023 22:29:58 +0300 Subject: [PATCH 1/6] added coef_net, prepare code fir nis_server --- integrator.py | 107 ++++++++++++++++++++++++++++---------------------- nis.py | 9 +++++ 2 files changed, 70 insertions(+), 46 deletions(-) diff --git a/integrator.py b/integrator.py index 990011d..3d54b60 100644 --- a/integrator.py +++ b/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 self.dist: torch.distributions.Distribution = dist self.optimizer = optimizer self.scheduler = scheduler @@ -252,7 +254,9 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False): ) with torch.no_grad(): x, absdet = self.flow(z, context=flow_context) + coef = self.coef_net(context[0, 1, 2, 6, 7]) # xyz , w0 return (x.to("cpu"), absdet.to("cpu")) + return (x.to("cpu"), absdet.to("cpu"), coef.to("cpu")) def train_with_context( self, @@ -284,62 +288,73 @@ def train_with_context( 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) 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_с) # 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) - # --------------- START TODO compute loss --------------- - start = time.time() - y = batch_y.to(self.device) + # cross = torch.nn.CrossEntropyLoss() + # loss = cross(absdet, y) + kl_loss = torch.nn.KLDivLoss() + loss = kl_loss(torch.log_softmax(absdet, -1), torch.softmax(y, -1)) - cross = torch.nn.CrossEntropyLoss() - loss = cross(absdet, y) - mean = torch.mean(y / absdet) - var = torch.var(y / absdet) + # absdet_coef = coef * absdet + (1 - coef) * hybrid_pdf + # loss_coef = kl_loss(torch.log_softmax(absdet_coef, -1), torch.softmax(y, -1)) + # + # 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 + 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) + # --------------- 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): diff --git a/nis.py b/nis.py index 13b9d68..9eeacf6 100644 --- a/nis.py +++ b/nis.py @@ -75,6 +75,14 @@ def initialize(self, mode="server"): for mask in masks ] ) + coef_net = UNet( + in_features=6, + 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 +96,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 = [] From 97bc84e3d33d3d2380d4898035ce10a08e396446 Mon Sep 17 00:00:00 2001 From: alexpetersburg Date: Sun, 12 Feb 2023 16:21:32 +0300 Subject: [PATCH 2/6] integrate mis into server_nis.py and client.py --- client.py | 8 +++++--- experiment_config.py | 4 ++-- integrator.py | 1 - nis.py | 4 ++-- server_nis.py | 24 ++++++++++++++---------- 5 files changed, 23 insertions(+), 18 deletions(-) diff --git a/client.py b/client.py index cbc128e..0c3879f 100644 --- a/client.py +++ b/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) diff --git a/experiment_config.py b/experiment_config.py index 2f4a2d2..f2f2549 100644 --- a/experiment_config.py +++ b/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 diff --git a/integrator.py b/integrator.py index 3d54b60..68feba9 100644 --- a/integrator.py +++ b/integrator.py @@ -255,7 +255,6 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False): with torch.no_grad(): x, absdet = self.flow(z, context=flow_context) coef = self.coef_net(context[0, 1, 2, 6, 7]) # xyz , w0 - return (x.to("cpu"), absdet.to("cpu")) return (x.to("cpu"), absdet.to("cpu"), coef.to("cpu")) def train_with_context( diff --git a/nis.py b/nis.py index 9eeacf6..b6a942d 100644 --- a/nis.py +++ b/nis.py @@ -233,9 +233,9 @@ def binary_list(inval, length): 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) + [samples, pdf, coef] = self.integrator.sample_with_context(context, inverse=False) pdf_light_sample = torch.ones(pdf.size()) - return [samples, pdf_light_sample, pdf] + return [samples, pdf_light_sample, pdf, coef] def train(self, context): z = torch.stack(self.integrator.generate_z_by_context(context)) diff --git a/server_nis.py b/server_nis.py index cfb300a..06f9f8a 100644 --- a/server_nis.py +++ b/server_nis.py @@ -93,12 +93,13 @@ def make_infer(self): ) # lights(vec3), pdfs # pdf_light_samples = utils.get_pdf_by_samples_uniform(points[:, 8:]) # [samples, pdfs] = utils.get_test_samples_uniform(points) # lights(vec3), pdfs + coef = torch.ones(pdf_light_samples.size()) else: if (self.nis.num_frame != 1) and ( not self.config.one_bounce_mode or (self.nis.train_sampling_call_difference == 1) ): - [samples, pdf_light_samples, pdfs] = self.nis.get_samples(points) + [samples, pdf_light_samples, pdfs, coef] = self.nis.get_samples(points) self.samples_tensor = ( samples.clone().numpy() ) # This is only needed for the make_train step and pass it to the Gaussian function. @@ -120,12 +121,13 @@ def make_infer(self): [samples, pdfs] = utils.get_test_samples_uniform( points ) # lights(vec3), pdfs + coef = torch.ones(pdf_light_samples.size()) - return [samples, pdf_light_samples, pdfs] + return [samples, pdf_light_samples, pdfs, coef] def make_train(self): # context = np.frombuffer(self.raw_data, dtype=np.float32).reshape((-1, self.config.num_context_features + 3)) - context = np.frombuffer(self.raw_data, dtype=np.float32).reshape((-1, 8 + 3)) + context = np.frombuffer(self.raw_data, dtype=np.float32).reshape((-1, 8 + 3 + 1)) context = context[~np.isnan(context).any(axis=1), :] if self.hybrid_sampling: pass @@ -136,11 +138,11 @@ def make_train(self): ): lum = 0.3 * context[:, 0] + 0.3 * context[:, 1] + 0.3 * context[:, 2] # Checking the Gaussian distribution - y = self.nis.function(torch.from_numpy(self.samples_tensor)) - lum[0] = y[0].item() - lum[1] = y[1].item() - lum[2] = y[2].item() - tdata = context[:, [3, 4, 5, 6, 7, 8, 9, 10]] + # y = self.nis.function(torch.from_numpy(self.samples_tensor)) + # lum[0] = y[0].item() + # lum[1] = y[1].item() + # lum[2] = y[2].item() + tdata = context[:, [3, 4, 5, 6, 7, 8, 9, 10, 11]] tdata = np.concatenate( (tdata, lum.reshape([len(lum), 1])), axis=1, dtype=np.float32 ) @@ -161,7 +163,7 @@ def process(self): self.make_train() self.connection.send(self.data_ok.name) elif self.mode == Mode.INFERENCE: - [samples, pdf_light_samples, pdfs] = self.make_infer() + [samples, pdf_light_samples, pdfs, coef] = self.make_infer() self.connection.send(self.put_infer.name) answer = self.connection.recv(self.put_infer_ok.length) if answer == self.put_infer_ok.name: @@ -173,7 +175,9 @@ def process(self): (s, pls.reshape([len(pls), 1])), axis=1, dtype=np.float32 ) p = pdfs.cpu().detach().numpy().reshape([-1, 1]) - raw_data.extend(np.concatenate((s, p), axis=1).tobytes()) + c = coef.cpu().detach().numpy().reshape([-1, 1]) + raw_data.extend(np.concatenate((s, p, c), axis=1).tobytes()) + self.connection.send(len(raw_data).to_bytes(4, "little")) self.connection.sendall(raw_data) answer = self.connection.recv(self.data_ok.length) From 3c9c6e3de877265304f56dd2c385729d34c57817 Mon Sep 17 00:00:00 2001 From: alexpetersburg Date: Tue, 28 Feb 2023 15:02:36 +0300 Subject: [PATCH 3/6] fix coef net input context, --- integrator.py | 2 +- tests/test_nis_client_server.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/integrator.py b/integrator.py index 68feba9..128a3af 100644 --- a/integrator.py +++ b/integrator.py @@ -254,7 +254,7 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False): ) with torch.no_grad(): x, absdet = self.flow(z, context=flow_context) - coef = self.coef_net(context[0, 1, 2, 6, 7]) # xyz , w0 + coef = self.coef_net(context[:, [0, 1, 2, 6, 7]], context=None) # xyz , w0 return (x.to("cpu"), absdet.to("cpu"), coef.to("cpu")) def train_with_context( diff --git a/tests/test_nis_client_server.py b/tests/test_nis_client_server.py index feab248..1ad6e3f 100644 --- a/tests/test_nis_client_server.py +++ b/tests/test_nis_client_server.py @@ -19,7 +19,7 @@ def test_nis_client_server(): def server_process_func(): - server = TrainServer(ExperimentConfig(num_context_features=6, save_plots=False)) + server = TrainServer(ExperimentConfig(num_context_features=8, save_plots=False)) server.nis.initialize(mode='server') server.connect() @@ -32,6 +32,8 @@ def server_process_func(): server.process() # inference nis + server.nis.num_frame = 2 + server.nis.train_sampling_call_difference = 0 server.hybrid_sampling = False server.mode = Mode.INFERENCE server.connection.send(server.put_infer_ok.name) From 1c14fd2094b5dd7af7dfce14e5934a1032cc93df Mon Sep 17 00:00:00 2001 From: alexpetersburg Date: Tue, 28 Feb 2023 15:07:10 +0300 Subject: [PATCH 4/6] fix coef net input context --- integrator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integrator.py b/integrator.py index 128a3af..5fb3052 100644 --- a/integrator.py +++ b/integrator.py @@ -254,7 +254,7 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False): ) with torch.no_grad(): x, absdet = self.flow(z, context=flow_context) - coef = self.coef_net(context[:, [0, 1, 2, 6, 7]], context=None) # xyz , w0 + coef = self.coef_net(torch.tensor(context[:, [0, 1, 2, 6, 7]]), context=None) # xyz , w0 return (x.to("cpu"), absdet.to("cpu"), coef.to("cpu")) def train_with_context( From d03912e80a9788efaa6dff6787133bf94b9fa538 Mon Sep 17 00:00:00 2001 From: alexpetersburg Date: Tue, 28 Feb 2023 15:18:33 +0300 Subject: [PATCH 5/6] fix coef net device --- integrator.py | 10 +++++----- nis.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/integrator.py b/integrator.py index 5fb3052..8e644f3 100644 --- a/integrator.py +++ b/integrator.py @@ -43,7 +43,7 @@ def __init__( self.global_step = 0 self.scheduler_step = 0 self.flow = flow.to(device) - self.coef_net = coef_net + self.coef_net = coef_net.to(device) self.dist: torch.distributions.Distribution = dist self.optimizer = optimizer self.scheduler = scheduler @@ -254,7 +254,7 @@ def sample_with_context(self, context: np.ndarray, inverse: bool = False): ) with torch.no_grad(): x, absdet = self.flow(z, context=flow_context) - coef = self.coef_net(torch.tensor(context[:, [0, 1, 2, 6, 7]]), context=None) # xyz , w0 + 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( @@ -276,7 +276,7 @@ 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) elif self.features_mode == "xyz": context_x = torch.Tensor(context[:, :3]).to(self.device) else: @@ -288,14 +288,14 @@ def train_with_context( start = time.time() context_y = torch.Tensor(context_y).to(self.device) - context_с = torch.Tensor(context[0, 1, 2, 6, 7]).to(self.device) + context_с = torch.Tensor(context[:, [0, 1, 2, 6, 7]]).to(self.device) logging.info(f"context_y time: {time.time() - start}") start = time.time() logging.info(f"batch time: {time.time() - start}") start = time.time() x, absdet = self.flow(z, context_x) - coef = self.coef_net(context_с) # xyz , w0 + coef = self.coef_net(context_с, context=None) # xyz , w0 absdet.requires_grad_(True) # y = self._func(x) logging.info(f"flow time: {time.time() - start}") diff --git a/nis.py b/nis.py index b6a942d..c59fc5b 100644 --- a/nis.py +++ b/nis.py @@ -76,7 +76,7 @@ def initialize(self, mode="server"): ] ) coef_net = UNet( - in_features=6, + in_features=5, out_features=1, max_hidden_features=self.config.hidden_dim, num_layers=self.config.n_hidden_layers, From bdc0e26e903ac5ad36d688f3b412e9c36f0a4bf8 Mon Sep 17 00:00:00 2001 From: alexpetersburg Date: Mon, 6 Mar 2023 23:55:46 +0300 Subject: [PATCH 6/6] added gradient_accumulation, drop_zero samples, experiments with loss functions --- experiment_config.py | 2 ++ integrator.py | 43 ++++++++++++++++++++---------- nis.py | 18 +++++++++---- server_nis.py | 63 +++++++++++++++++++++++--------------------- 4 files changed, 77 insertions(+), 49 deletions(-) diff --git a/experiment_config.py b/experiment_config.py index f2f2549..693fbef 100644 --- a/experiment_config.py +++ b/experiment_config.py @@ -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"), diff --git a/integrator.py b/integrator.py index 8e644f3..9aefbac 100644 --- a/integrator.py +++ b/integrator.py @@ -229,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 @@ -241,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: @@ -277,11 +280,12 @@ def train_with_context( # Process # if self.features_mode == "all_features": 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}") @@ -289,13 +293,14 @@ def train_with_context( 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() logging.info(f"batch time: {time.time() - start}") start = time.time() x, absdet = self.flow(z, context_x) - coef = self.coef_net(context_с, context=None) # xyz , w0 + 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}") @@ -304,18 +309,27 @@ def train_with_context( # --------------- START TODO compute loss --------------- start = time.time() y = context_y.to(self.device) - - # cross = torch.nn.CrossEntropyLoss() - # loss = cross(absdet, y) + y = y + 1e-10 + cross = torch.nn.CrossEntropyLoss() + loss = cross(absdet, y) kl_loss = torch.nn.KLDivLoss() - loss = kl_loss(torch.log_softmax(absdet, -1), torch.softmax(y, -1)) + 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() - # absdet_coef = coef * absdet + (1 - coef) * hybrid_pdf + absdet_coef = coef * absdet + (1 - coef) * hybrid_pdf # loss_coef = kl_loss(torch.log_softmax(absdet_coef, -1), torch.softmax(y, -1)) - # - # tao = 0.5 - # b_tao = 0.5*((1/3)**(5*tao)) - # final_loss = b_tao*loss + (1-b_tao)*loss_coef + loss_coef = (y*torch.log(y)).mean() - (y * torch.log(absdet_coef)).mean() + + + 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) @@ -329,9 +343,10 @@ def train_with_context( logging.info(f"loss time: {time.time() - start}") start = time.time() - loss.backward() + # 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: diff --git a/nis.py b/nis.py index c59fc5b..1e224b0 100644 --- a/nis.py +++ b/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: @@ -232,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) + 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()) + # 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: @@ -256,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))[ @@ -270,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}") @@ -284,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 diff --git a/server_nis.py b/server_nis.py index 06f9f8a..0271f42 100644 --- a/server_nis.py +++ b/server_nis.py @@ -87,43 +87,46 @@ def make_infer(self): (-1, 8 + 2) ) # Temporal solution for ignoring processing additional inputs in NN if self.hybrid_sampling: - pdf_light_samples = utils.get_pdf_by_samples_cosine(points[:, 8:]) - [samples, pdfs] = utils.get_test_samples_cosine( - points - ) # lights(vec3), pdfs - # pdf_light_samples = utils.get_pdf_by_samples_uniform(points[:, 8:]) - # [samples, pdfs] = utils.get_test_samples_uniform(points) # lights(vec3), pdfs - coef = torch.ones(pdf_light_samples.size()) + samples, pdfs, pdf_light_samples, coef = self._make_infer_hybrid(points) else: if (self.nis.num_frame != 1) and ( not self.config.one_bounce_mode or (self.nis.train_sampling_call_difference == 1) ): - [samples, pdf_light_samples, pdfs, coef] = self.nis.get_samples(points) - self.samples_tensor = ( - samples.clone().numpy() - ) # This is only needed for the make_train step and pass it to the Gaussian function. - samples[:, 0] = samples[:, 0] * 2 * np.pi - samples[:, 1] = torch.acos(samples[:, 1]) - - # MIS should be implemented here - # pdfs = (1 / (2 * np.pi)) / pdfs - # pdf_light_samples = pdf_light_samples / (2 * np.pi) - # if self.nis.num_frame < 100: - # pdfs = torch.ones(pdfs.size()) - # pdfs /= (2 * np.pi) - # pdf_light_samples = torch.ones(pdfs.size()) - # pdf_light_samples /= (2 * np.pi) + samples, pdfs, pdf_light_samples, coef = self._make_infer_nis(points) else: - # pdf_light_samples = utils.get_pdf_by_samples_cosine(points[:, 8:]) - # [samples, pdfs] = utils.get_test_samples_cosine(points) # lights(vec3), pdfs - pdf_light_samples = utils.get_pdf_by_samples_uniform(points[:, 8:]) - [samples, pdfs] = utils.get_test_samples_uniform( - points - ) # lights(vec3), pdfs - coef = torch.ones(pdf_light_samples.size()) - + samples, pdfs, pdf_light_samples, coef = self._make_infer_hybrid(points) return [samples, pdf_light_samples, pdfs, coef] + # return [samples, pdf_light_samples, torch.nn.Softmax()(pdfs), coef] + + def _make_infer_hybrid(self, points): + pdf_light_samples = utils.get_pdf_by_samples_cosine(points[:, 8:]) + # [samples, pdfs] = utils.get_test_samples_cosine( + # points + # ) + [samples, pdfs] = utils.get_test_samples_vectorized(points)# lights(vec3), pdfs + # pdf_light_samples = utils.get_pdf_by_samples_uniform(points[:, 8:]) + # [samples, pdfs] = utils.get_test_samples_uniform(points) # lights(vec3), pdfs + coef = torch.ones(pdf_light_samples.size()) + return samples, pdfs, pdf_light_samples, coef + + def _make_infer_nis(self, points): + [samples, pdf_light_samples, pdfs, coef] = self.nis.get_samples(points) + self.samples_tensor = ( + samples.clone().numpy() + ) # This is only needed for the make_train step and pass it to the Gaussian function. + samples[:, 0] = samples[:, 0] * 2 * np.pi + samples[:, 1] = torch.acos(samples[:, 1]) + + # MIS should be implemented here + # pdfs = (1 / (2 * np.pi)) / pdfs + # pdf_light_samples = pdf_light_samples / (2 * np.pi) + # if self.nis.num_frame < 100: + # pdfs = torch.ones(pdfs.size()) + # pdfs /= (2 * np.pi) + # pdf_light_samples = torch.ones(pdfs.size()) + # pdf_light_samples /= (2 * np.pi) + return samples, pdfs, pdf_light_samples, coef def make_train(self): # context = np.frombuffer(self.raw_data, dtype=np.float32).reshape((-1, self.config.num_context_features + 3))