Spaces:
Runtime error
Runtime error
| """ Moderator | |
| # Input feature: body, part(head, hand) | |
| # output: fused feature, weight | |
| """ | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| # MLP + temperature softmax | |
| # w = SoftMax(w^\prime * temperature) | |
| class TempSoftmaxFusion(nn.Module): | |
| def __init__(self, channels=[2048 * 2, 1024, 1], detach_inputs=False, detach_feature=False): | |
| super(TempSoftmaxFusion, self).__init__() | |
| self.detach_inputs = detach_inputs | |
| self.detach_feature = detach_feature | |
| # weight | |
| layers = [] | |
| for l in range(0, len(channels) - 1): | |
| layers.append(nn.Linear(channels[l], channels[l + 1])) | |
| if l < len(channels) - 2: | |
| layers.append(nn.ReLU()) | |
| self.layers = nn.Sequential(*layers) | |
| # temperature | |
| self.register_parameter("temperature", nn.Parameter(torch.ones(1))) | |
| def forward(self, x, y, work=True): | |
| """ | |
| x: feature from body | |
| y: feature from part(head/hand) | |
| work: whether to fuse features | |
| """ | |
| if work: | |
| # 1. cat input feature, predict the weights | |
| f_in = torch.cat([x, y], dim=1) | |
| if self.detach_inputs: | |
| f_in = f_in.detach() | |
| f_temp = self.layers(f_in) | |
| f_weight = F.softmax(f_temp * self.temperature, dim=1) | |
| # 2. feature fusion | |
| if self.detach_feature: | |
| x = x.detach() | |
| y = y.detach() | |
| f_out = f_weight[:, [0]] * x + f_weight[:, [1]] * y | |
| x_out = f_out | |
| y_out = f_out | |
| else: | |
| x_out = x | |
| y_out = y | |
| f_weight = None | |
| return x_out, y_out, f_weight | |
| # MLP + Gumbel-Softmax trick | |
| # w = w^{\prime} - w^{\prime}\text{.detach()} + w^{\prime}\text{.gt(0.5)} | |
| class GumbelSoftmaxFusion(nn.Module): | |
| def __init__(self, channels=[2048 * 2, 1024, 1], detach_inputs=False, detach_feature=False): | |
| super(GumbelSoftmaxFusion, self).__init__() | |
| self.detach_inputs = detach_inputs | |
| self.detach_feature = detach_feature | |
| # weight | |
| layers = [] | |
| for l in range(0, len(channels) - 1): | |
| layers.append(nn.Linear(channels[l], channels[l + 1])) | |
| if l < len(channels) - 2: | |
| layers.append(nn.ReLU()) | |
| layers.append(nn.Softmax()) | |
| self.layers = nn.Sequential(*layers) | |
| def forward(self, x, y, work=True): | |
| """ | |
| x: feature from body | |
| y: feature from part(head/hand) | |
| work: whether to fuse features | |
| """ | |
| if work: | |
| # 1. cat input feature, predict the weights | |
| f_in = torch.cat([x, y], dim=-1) | |
| if self.detach_inputs: | |
| f_in = f_in.detach() | |
| f_weight = self.layers(f_in) | |
| # weight to be hard | |
| f_weight = f_weight - f_weight.detach() + f_weight.gt(0.5) | |
| # 2. feature fusion | |
| if self.detach_feature: | |
| x = x.detach() | |
| y = y.detach() | |
| f_out = f_weight[:, [0]] * x + f_weight[:, [1]] * y | |
| x_out = f_out | |
| y_out = f_out | |
| else: | |
| x_out = x | |
| y_out = y | |
| f_weight = None | |
| return x_out, y_out, f_weight | |