| import torch |
| import torch.nn as nn |
| from torch.nn import functional as F |
|
|
|
|
| class SemanticEmbedding(nn.Module): |
| def __init__(self, args, mesh_dim=71, report_dim=761, embed_size=512): |
| super(SemanticEmbedding, self).__init__() |
| self.mesh_tf = nn.Sequential( |
| nn.Linear(embed_size, embed_size // 2), |
| nn.ReLU(inplace=True), |
| nn.Linear(embed_size // 2, embed_size // 4), |
| nn.ReLU(inplace=True), |
| nn.Linear(embed_size // 4, mesh_dim) |
| ) |
|
|
| self.report_tf = nn.Sequential( |
| nn.Linear(embed_size, embed_size // 2), |
| nn.ReLU(inplace=True), |
| nn.Linear(embed_size // 2, embed_size // 4), |
| nn.ReLU(inplace=True), |
| nn.Linear(embed_size // 4, report_dim) |
| ) |
| self.bn = nn.BatchNorm1d(num_features=embed_size, momentum=0.1) |
| self.w1 = nn.Linear(in_features=mesh_dim + report_dim, out_features=embed_size) |
| self.w2 = nn.Linear(in_features=embed_size, out_features=embed_size) |
| self.relu = nn.ReLU() |
| self.logit = nn.Linear(60, 31) |
| self.dropout = nn.Dropout(0.2) |
| self.__init_weight() |
| self.target_dim = 60 |
| self.sigm = nn.Sigmoid() |
|
|
| def __init_weight(self): |
| self.w1.weight.data.uniform_(-0.1, 0.1) |
| self.w1.bias.data.fill_(0) |
| self.w2.weight.data.uniform_(-0.1, 0.1) |
| self.w2.bias.data.fill_(0) |
|
|
| def forward(self, avg, pred_output): |
| avg_visual = avg.unsqueeze(1) |
| pred_output2 = F.pad(pred_output, (0, 0, 0, self.target_dim - pred_output.shape[1]), 'constant', 0) |
| pred = pred_output2.permute(0, 2, 1) |
| visual_text = torch.matmul(avg_visual, pred).squeeze(1) |
| outputs = self.sigm(self.logit(visual_text)) |
| return outputs |
|
|
|
|
| class classfication(nn.Module): |
| def __init__(self, distiller_num, avg_dim=1024): |
| super(classfication, self).__init__() |
| self.logit = nn.Linear(avg_dim, distiller_num) |
| self.relu = nn.ReLU() |
| self.sigm = nn.Sigmoid() |
| self.dropout = nn.Dropout(0.5) |
|
|
| def forward(self, avg): |
| avg_visual = self.dropout(avg) |
| x = self.logit(avg_visual) |
| outputs = self.sigm(x) |
| return outputs |
|
|
