def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio, aligned): ctx.save_for_backward(roi) ctx.output_size = _pair(output_size) ctx.spatial_scale = spatial_scale ctx.sampling_ratio = sampling_ratio ctx.input_shape = input.size() ctx.aligned = aligned output = _C.roi_align_forward( input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio, aligned ) return output
def bench(): _C.roi_align_forward(input, boxes, 1.0, 7, 7, 0, True) torch.cuda.synchronize()