Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from deepsvg.difflib.tensor import SVGTensor | |
| from .utils import _get_padding_mask, _get_visibility_mask | |
| from .config import _DefaultConfig | |
| class SVGLoss(nn.Module): | |
| def __init__(self, cfg: _DefaultConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.args_dim = 2 * cfg.args_dim if cfg.rel_targets else cfg.args_dim + 1 | |
| self.register_buffer("cmd_args_mask", SVGTensor.CMD_ARGS_MASK) | |
| def forward(self, output, labels, weights): | |
| loss = 0. | |
| res = {} | |
| # VAE | |
| if self.cfg.use_vae: | |
| mu, logsigma = output["mu"], output["logsigma"] | |
| loss_kl = -0.5 * torch.mean(1 + logsigma - mu.pow(2) - torch.exp(logsigma)) | |
| loss_kl = loss_kl.clamp(min=weights["kl_tolerance"]) | |
| loss += weights["loss_kl_weight"] * loss_kl | |
| res["loss_kl"] = loss_kl | |
| # remove commitment loss | |
| # if self.cfg.use_vqvae: | |
| # vqvae_loss = output["vqvae_loss"].mean() | |
| # loss += vqvae_loss | |
| # res["vqvae_loss"] = vqvae_loss | |
| # Target & predictions | |
| # tgt_commands.shape [batch_size, max_num_groups, max_seq_len + 2] | |
| # tgt_args.shape [batch_size, max_num_groups, max_seq_len + 2, n_args] | |
| tgt_commands, tgt_args = output["tgt_commands"], output["tgt_args"] | |
| visibility_mask = _get_visibility_mask(tgt_commands, seq_dim=-1) | |
| padding_mask = _get_padding_mask(tgt_commands, seq_dim=-1, extended=True) * visibility_mask.unsqueeze(-1) | |
| command_logits, args_logits = output["command_logits"], output["args_logits"] | |
| # 2-stage visibility | |
| if self.cfg.decode_stages == 2: | |
| visibility_logits = output["visibility_logits"] | |
| loss_visibility = F.cross_entropy(visibility_logits.reshape(-1, 2), visibility_mask.reshape(-1).long()) | |
| loss += weights["loss_visibility_weight"] * loss_visibility | |
| res["loss_visibility"] = loss_visibility | |
| # Commands & args | |
| if self.cfg.bin_targets: # 当使用 bin_targets 时，每个坐标是由 8 bit 代表的，所以会多一维 | |
| tgt_args = tgt_args[..., 1:, :, :] | |
| else: | |
| tgt_args = tgt_args[..., 1:, :] | |
| tgt_commands, padding_mask = tgt_commands[..., 1:], padding_mask[..., 1:] | |
| # mask.shape [batch_size, 8, 31, 11] | |
| # 对于预测正确的 command, mask 会乘上 True, cmd_args_mask 向量不会发生改变 | |
| # 对于预测错误的 command, mask 会乘上 False, 相当于把 cmd_args_mask 置为 0, 即不统计对应的 args | |
| # pred_cmd = torch.argmax(command_logits, dim = -1) | |
| # mask = self.cmd_args_mask[tgt_commands.long()] * (pred_cmd == tgt_commands).unsqueeze(-1) | |
| mask = self.cmd_args_mask[tgt_commands.long()] | |
| # padding_mask.shape [batch_size, num_path, num_commands + 1] | |
| # command_logits.shape [batch_size, num_path, num_commands + 1, n_commands] | |
| # command_logits[padding_mask.bool()].shape [-1, n_commands] | |
| # 目的是把 PAD 的位置筛掉 | |
| loss_cmd = F.cross_entropy(command_logits[padding_mask.bool()].reshape(-1, self.cfg.n_commands), tgt_commands[padding_mask.bool()].reshape(-1).long()) | |
| if self.cfg.abs_targets: | |
| # l2 loss performs better than l1 loss | |
| loss_args = nn.MSELoss()( | |
| args_logits[mask.bool()].reshape(-1), | |
| tgt_args[mask.bool()].reshape(-1).float() | |
| ) | |
| elif self.cfg.bin_targets: | |
| loss_args = nn.MSELoss()( | |
| args_logits[mask.bool()].reshape(-1), | |
| tgt_args[mask.bool()].reshape(-1).float() | |
| ) | |
| else: | |
| loss_args = F.cross_entropy( | |
| args_logits[mask.bool()].reshape(-1, self.args_dim), | |
| tgt_args[mask.bool()].reshape(-1).long() + 1 | |
| ) # shift due to -1 PAD_VAL | |
| loss += weights["loss_cmd_weight"] * loss_cmd \ | |
| + weights["loss_args_weight"] * loss_args | |
| res.update({ | |
| "loss": loss, | |
| "loss_cmd": loss_cmd, | |
| "loss_args": loss_args | |
| }) | |
| return res | |