Putting LoadLibraryA("torch_cuda.dll"); early in my code allowed me to start using the nightly debug build of libtorch 1.5, but I’m still stuck on module.forward.
2020-07-26 21_27_53-Microsoft Visual Studio1163×384 26.4 KB
I also put -INCLUDE:?warp_size@cuda@at@@YAHXZ in Linker>All Options>Additional Options.
and LoadLibraryA("c10_cuda.dll"); early too. Is there something else I can try?
Update again: I can trace a style transfer model and that works in my code here, but the traced U2Net model doesn’t work.
Below is how I wrapped the model. Is it ok to use x[:,0,:,:] or does that break the jit trace? I’m also concerned about this line https://github.com/NathanUA/U-2-Net/blob/b77cd6da3204efcb03e18e15dd3b9eb24d47f969/model/u2net.py#L24
def normPRED(d):
    ma = torch.max(d)
    mi = torch.min(d)

    dn = (d-mi)/(ma-mi)

    return dn


class ModelWrapper(nn.Module):

    def __init__(self, u2netmodel):
            super(ModelWrapper,self).__init__()

            self.u2netmodel = u2netmodel

    def forward(self, x):

            # my code doesn't use ToTensorLab in the data loader
            # https://github.com/NathanUA/U-2-Net/blob/b77cd6da3204efcb03e18e15dd3b9eb24d47f969/data_loader.py#L208
            # so do the normalization in this wrapper

            x = x / torch.max(x)

            r = (x[:,0,:,:]-0.485)/0.229
            g = (x[:,1,:,:]-0.456)/0.224
            b = (x[:,2,:,:]-0.406)/0.225

            img = torch.stack((r,g,b), 1)

            d1,d2,d3,d4,d5,d6,d7= self.u2netmodel(img)

            return normPRED(d1)

# paraphrasing u2net_test.py
net = U2NET(3,1)
net.cuda()
wrapper = ModelWrapper(net)
wrapper.cuda()
wrapper.to('cuda') # just in case?
print('is cuda: ' + str(next(wrapper.parameters()).is_cuda)) # True

inputs_test = data_test['image']
inputs_test = inputs_test.type(torch.FloatTensor)
inputs_test = inputs_test.cuda()
print("inputs size: " + str(inputs_test.size())) # [1, 3, 320, 320]

d1 = wrapper(inputs_test)

# save d1 as an image and the image is great!
# save_output(img_name_list[i_test],pred,prediction_dir)

traced = torch.jit.trace(wrapper, inputs_test)
traced.save("traced_model.pt")