def test(self):
"""tests model
returns: number correct and total number
"""
with torch.no_grad():
out = self.forward()
prediction, confidence = out
if self.opt.arch == "vae":
predicted_cp = utils.transform_control_points(
prediction, prediction.shape[0], device=self.device)
reconstruction_loss, _ = self.criterion[1](
predicted_cp,
self.targets,
confidence=confidence,
confidence_weight=self.opt.confidence_weight,
device=self.device)
return reconstruction_loss, 1
elif self.opt.arch == "gan":
predicted_cp = utils.transform_control_points(
prediction, prediction.shape[0], device=self.device)
reconstruction_loss, _ = self.criterion(
predicted_cp,
self.targets,
confidence=confidence,
confidence_weight=self.opt.confidence_weight,
device=self.device)
return reconstruction_loss, 1
else:
predicted = torch.round(torch.sigmoid(prediction)).squeeze()
correct = (predicted == self.targets).sum().item()
return correct, len(self.targets)
def backward(self, out):
if self.opt.arch == 'vae':
predicted_cp, confidence, mu, logvar = out
predicted_cp = utils.transform_control_points(
predicted_cp, predicted_cp.shape[0], device=self.device)
self.reconstruction_loss, self.confidence_loss = self.criterion[1](
predicted_cp,
self.targets,
confidence=confidence,
confidence_weight=self.opt.confidence_weight,
device=self.device)
self.kl_loss = self.opt.kl_loss_weight * self.criterion[0](
mu, logvar, device=self.device)
self.l2_loss = self.criterion[2](
predicted_cp,
self.targets,
confidence=confidence,
confidence_weight=self.opt.confidence_weight,
device=self.device)
self.trans_l2_loss = self.criterion[3](
predicted_cp,
self.targets,
confidence=confidence,
confidence_weight=self.opt.confidence_weight,
device=self.device)
self.loss = self.kl_loss + self.reconstruction_loss + self.confidence_loss + 0.2*self.l2_loss + 0.8*self.trans_l2_loss
elif self.opt.arch == 'gan':
predicted_cp, confidence = out
predicted_cp = utils.transform_control_points(
predicted_cp, predicted_cp.shape[0], device=self.device)
self.reconstruction_loss, self.confidence_loss = self.criterion(
predicted_cp,
self.targets,
confidence=confidence,
confidence_weight=self.opt.confidence_weight,
device=self.device)
self.loss = self.reconstruction_loss + self.confidence_loss
elif self.opt.arch == 'evaluator':
grasp_classification, confidence = out
self.classification_loss, self.confidence_loss = self.criterion(
grasp_classification.squeeze(),
self.targets,
confidence,
self.opt.confidence_weight,
device=self.device)
self.loss = self.classification_loss + self.confidence_loss
self.loss.backward()