|
- # Copyright (c) Microsoft Corporation.
- # Licensed under the MIT license.
-
- import logging
-
- import torch
- import torch.nn as nn
- import json
-
- INF = 1E10
- EPS = 1E-12
-
- logger = logging.getLogger(__name__)
- logger.setLevel(logging.INFO)
-
- def get_length(mask):
- length = torch.sum(mask, 1)
- length = length.long().cpu()
- return length
-
-
- class GlobalAvgPool(nn.Module):
- def forward(self, x, mask):
- x = torch.sum(x, 2)
- length = torch.sum(mask, 1, keepdim=True).float()
- length += torch.eq(length, 0.0).float() * EPS
- length = length.repeat(1, x.size()[1])
- x /= length
- return x
-
-
- class GlobalMaxPool(nn.Module):
- def forward(self, x, mask):
- mask = torch.eq(mask.float(), 0.0).long()
- mask = torch.unsqueeze(mask, dim=1).repeat(1, x.size()[1], 1)
- mask *= -INF
- x += mask
- x, _ = torch.max(x + mask, 2)
- return x
-
-
- class IteratorWrapper:
- def __init__(self, loader):
- self.loader = loader
- self.iterator = None
-
- def __iter__(self):
- self.iterator = iter(self.loader)
- return self
-
- def __len__(self):
- return len(self.loader)
-
- def __next__(self):
- data = next(self.iterator)
- text, length = data.text
- max_length = text.size(1)
- label = data.label - 1
- bs = label.size(0)
- mask = torch.arange(max_length, device=length.device).unsqueeze(0).repeat(bs, 1)
- mask = mask < length.unsqueeze(-1).repeat(1, max_length)
- return (text, mask), label
-
-
- def accuracy(output, target):
- batch_size = target.size(0)
- _, predicted = torch.max(output.data, 1)
- return (predicted == target).sum().item() / batch_size
-
-
- def dump_global_result(res_path,global_result, sort_keys = False):
- with open(res_path, "w") as ss_file:
- json.dump(global_result, ss_file, sort_keys=sort_keys, indent=2)
|
