_meta_registrations.py 1.59 KB
Newer Older
1
2
import functools

Edward Z. Yang's avatar
Edward Z. Yang committed
3
4
5
6
7
8
9
10
11
import torch
import torch.library

# Ensure that torch.ops.torchvision is visible
import torchvision.extension  # noqa: F401

from torch._prims_common import check


12
13
14
@functools.lru_cache(None)
def get_meta_lib():
    return torch.library.Library("torchvision", "IMPL", "Meta")
Edward Z. Yang's avatar
Edward Z. Yang committed
15
16


17
def register_meta(op_name, overload_name="default"):
Edward Z. Yang's avatar
Edward Z. Yang committed
18
    def wrapper(fn):
19
20
        if torchvision.extension._has_ops():
            get_meta_lib().impl(getattr(getattr(torch.ops.torchvision, op_name), overload_name), fn)
Edward Z. Yang's avatar
Edward Z. Yang committed
21
22
23
24
25
        return fn

    return wrapper


26
@register_meta("roi_align")
Edward Z. Yang's avatar
Edward Z. Yang committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def meta_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned):
    check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]")
    check(
        input.dtype == rois.dtype,
        lambda: (
            "Expected tensor for input to have the same type as tensor for rois; "
            f"but type {input.dtype} does not equal {rois.dtype}"
        ),
    )
    num_rois = rois.size(0)
    _, channels, height, width = input.size()
    return input.new_empty((num_rois, channels, pooled_height, pooled_width))


41
@register_meta("_roi_align_backward")
Edward Z. Yang's avatar
Edward Z. Yang committed
42
43
44
45
46
47
48
49
50
51
52
def meta_roi_align_backward(
    grad, rois, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width, sampling_ratio, aligned
):
    check(
        grad.dtype == rois.dtype,
        lambda: (
            "Expected tensor for grad to have the same type as tensor for rois; "
            f"but type {grad.dtype} does not equal {rois.dtype}"
        ),
    )
    return grad.new_empty((batch_size, channels, height, width))