| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def dice(x, y): |
| intersect = np.sum(np.sum(np.sum(x * y))) |
| y_sum = np.sum(np.sum(np.sum(y))) |
| if y_sum == 0: |
| return 0.0 |
| x_sum = np.sum(np.sum(np.sum(x))) |
| return 2 * intersect / (x_sum + y_sum) |
|
|
|
|
| class AverageMeter(object): |
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.val = 0 |
| self.avg = 0 |
| self.sum = 0 |
| self.count = 0 |
|
|
| def update(self, val, n=1): |
| self.val = val |
| self.sum += val * n |
| self.count += n |
| self.avg = np.where(self.count > 0, self.sum / self.count, self.sum) |
|
|
|
|
| def distributed_all_gather( |
| tensor_list, valid_batch_size=None, out_numpy=False, world_size=None, no_barrier=False, is_valid=None |
| ): |
| if world_size is None: |
| world_size = torch.distributed.get_world_size() |
| if valid_batch_size is not None: |
| valid_batch_size = min(valid_batch_size, world_size) |
| elif is_valid is not None: |
| is_valid = torch.tensor(bool(is_valid), dtype=torch.bool, device=tensor_list[0].device) |
| if not no_barrier: |
| torch.distributed.barrier() |
| tensor_list_out = [] |
| with torch.no_grad(): |
| if is_valid is not None: |
| is_valid_list = [torch.zeros_like(is_valid) for _ in range(world_size)] |
| torch.distributed.all_gather(is_valid_list, is_valid) |
| is_valid = [x.item() for x in is_valid_list] |
| for tensor in tensor_list: |
| gather_list = [torch.zeros_like(tensor) for _ in range(world_size)] |
| torch.distributed.all_gather(gather_list, tensor) |
| if valid_batch_size is not None: |
| gather_list = gather_list[:valid_batch_size] |
| elif is_valid is not None: |
| gather_list = [g for g, v in zip(gather_list, is_valid_list) if v] |
| if out_numpy: |
| gather_list = [t.cpu().numpy() for t in gather_list] |
| tensor_list_out.append(gather_list) |
| return tensor_list_out |
|
|
