|
- """
- # -*- coding: utf-8 -*-
- -----------------------------------------------------------------------------------
- # Author: Nguyen Mau Dung
- # DoC: 2020.08.09
- # email: nguyenmaudung93.kstn@gmail.com
- -----------------------------------------------------------------------------------
- # Description: some utilities of torch (conversion)
- -----------------------------------------------------------------------------------
- """
-
- import torch
- import torch.distributed as dist
-
- __all__ = ['convert2cpu', 'convert2cpu_long', 'to_cpu', 'reduce_tensor', 'to_python_float', '_sigmoid']
-
-
- def convert2cpu(gpu_matrix):
- return torch.FloatTensor(gpu_matrix.size()).copy_(gpu_matrix)
-
-
- def convert2cpu_long(gpu_matrix):
- return torch.LongTensor(gpu_matrix.size()).copy_(gpu_matrix)
-
-
- def to_cpu(tensor):
- return tensor.detach().cpu()
-
-
- def reduce_tensor(tensor, world_size):
- rt = tensor.clone()
- dist.all_reduce(rt, op=dist.reduce_op.SUM)
- rt /= world_size
- return rt
-
-
- def to_python_float(t):
- if hasattr(t, 'item'):
- return t.item()
- else:
- return t[0]
-
-
- def _sigmoid(x):
- return torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4)
|
