| 723 | |
| 724 | |
| 725 | class Aggregation(Function): |
| 726 | @staticmethod |
| 727 | def forward(ctx, input, position, weight, idx): |
| 728 | """ |
| 729 | input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample) |
| 730 | output: (n, c) |
| 731 | """ |
| 732 | assert input.is_contiguous() and position.is_contiguous() and weight.is_contiguous() |
| 733 | n, nsample, c = position.shape; w_c = weight.shape[-1] |
| 734 | output = torch.cuda.FloatTensor(n, c).zero_() |
| 735 | pointops_cuda.aggregation_forward_cuda(n, nsample, c, w_c, input, position, weight, idx, output) |
| 736 | ctx.save_for_backward(input, position, weight, idx) |
| 737 | return output |
| 738 | |
| 739 | @staticmethod |
| 740 | def backward(ctx, grad_output): |
| 741 | """ |
| 742 | input: grad_out: (n, c) |
| 743 | output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c') |
| 744 | """ |
| 745 | input, position, weight, idx = ctx.saved_tensors |
| 746 | n, nsample, c = position.shape; w_c = weight.shape[-1] |
| 747 | grad_input = torch.cuda.FloatTensor(n, c).zero_() |
| 748 | grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_() |
| 749 | grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_() |
| 750 | pointops_cuda.aggregation_backward_cuda(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight) |
| 751 | return grad_input, grad_position, grad_weight, None |
| 752 | |
| 753 | aggregation = Aggregation.apply |
| 754 |
nothing calls this directly
no outgoing calls
no test coverage detected