def training_step(self, batch, _batch_nb):
# TODO: plot
y, x, uids = (emiss, laser_params, uids) = batch
x_pred = self(y)
with torch.no_grad():
x_loss = rmse(x_pred, x)
self.log("backward/train/x/loss", x_loss, prog_bar=True)
if self.forward_model is not None:
y_pred = self.forward_model(x_pred)
y_loss = rmse(y_pred, y)
self.log(
"backward/train/y/loss",
y_loss,
prog_bar=True,
)
loss = y_loss
if self.current_epoch == self.config["backward_num_epochs"] - 5:
nngraph.save_integral_emiss_point(
y_pred,
y,
"/data-new/alok/laser/backwards_train_points.txt",
all_points=True,
)
self.log(f"backward/train/loss", loss, prog_bar=True)
return loss
def test_step(self, batch, batch_nb):
# TODO: plot
y, x, uids = (emiss, laser_params, uids) = batch
x_pred = self(y)
with torch.no_grad():
x_loss = rmse(x_pred, x)
self.log("backward/test/x/loss", x_loss, prog_bar=True)
if self.forward_model is not None:
y_pred = self.forward_model(x_pred)
y_loss = rmse(y_pred, y)
self.log(
"backward/test/y/loss",
y_loss,
prog_bar=True,
)
loss = y_loss
torch.save(x, "/data-new/alok/laser/params_true_back.pt")
torch.save(y, "/data-new/alok/laser/emiss_true_back.pt")
torch.save(y_pred, "/data-new/alok/laser/emiss_pred.pt")
torch.save(x_pred, "/data-new/alok/laser/param_pred.pt")
nngraph.save_integral_emiss_point(
y_pred,
y,
"/data-new/alok/laser/backwards_test_points.txt",
all_points=True)
return loss
def test_step(self, batch, batch_nb):
x, y, uids = batch
y_pred = self(x)
loss = rmse(y_pred, y)
self.log(f"forward/test/loss", loss, prog_bar=True)
nngraph.save_integral_emiss_point(
y_pred, y, "/data-new/alok/laser/forwards_val_points.txt", all_points=True
)
return loss
def validation_step(self, batch, batch_nb):
x, y, uids = batch
y_pred = self(x)
loss = rmse(y_pred, y)
randcheck = np.random.uniform()
self.log(f"forward/val/loss", loss, prog_bar=True)
if self.current_epoch > self.config["forward_num_epochs"] - 5:
nngraph.save_integral_emiss_point(
y_pred, y, "/data-new/alok/laser/forwards_val_points.txt", all_points=True
)
return loss
def training_step(self, batch, batch_nb):
x, y, uids = batch
y_pred = self(x)
loss = rmse(y_pred, y)
# nngraph.emiss_error_graph(y_pred, y, "train_step.png")
# self.log_image(key="train_forwards_error_graphs", images=["train_step.png"])
if self.current_epoch == self.config["forward_num_epochs"] - 5:
nngraph.save_integral_emiss_point(
y_pred, y, "/data-new/alok/laser/forwards_train_points.txt", all_points=True
)
self.log(f"forward/train/loss", loss, prog_bar=True)
return loss