wkcn / MobulaOP

A Simple & Flexible Cross Framework Operators Toolkit
MIT License
164 stars 21 forks source link

GPU backend does not work for pytorch #9

Closed merrymercy closed 5 years ago

merrymercy commented 6 years ago

The following code produces wrong output. If I change .cuda() to .cpu(), I can get correct output.

(Fix #10 is required to run this example)

# Use ROIAlign operator
import sys
sys.path.append('../') # Add MobulaOP path
import numpy as np
import mobula
# Load ROIAlign Module
mobula.op.load('ROIAlign')

dtype = np.float32
N, C, H, W = 2, 3, 4, 4

import torch

data = torch.tensor(np.arange(N*C*H*W).astype(dtype).reshape((N,C,H,W))).cuda()
rois = torch.tensor(np.array([[0, 1, 1, 3, 3]], dtype = dtype)).cuda()

output = mobula.op.ROIAlign(data = data, rois = rois, pooled_size = (2,2), spatial_scale = 1.0, sampling_ratio = 1)

print("= OUTPUT =")
print (output)
merrymercy commented 6 years ago

11 fixes this issue partially

wkcn commented 6 years ago

Thank you very much! I have merged your two PRs.