CodeMirror/waifu2x
1
0
Fork 0
mirror of synced 2024-06-13 08:24:30 +12:00
Code Issues Releases Wiki Activity

Fix cuda tensor error in async environment

Browse source
This commit is contained in:
nagadomi 2015-12-01 18:21:06 +09:00
parent d2c081bbcf
commit c5df72439f
1 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
lib/alpha_util.lua
View file

@ -5,15 +5,15 @@ local iproc = require 'iproc'
local gm = require 'graphicsmagick'
alpha_util = {}
alpha_util.sum2d = nn.SpatialConvolutionMM(1, 1, 3, 3, 1, 1, 1, 1):cuda()
alpha_util.sum2d.weight:fill(1)
alpha_util.sum2d.bias:zero()
function alpha_util.make_border(rgb, alpha, offset)
if not alpha then
return rgb
end
local sum2d = nn.SpatialConvolutionMM(1, 1, 3, 3, 1, 1, 1, 1):cuda()
sum2d.weight:fill(1)
sum2d.bias:zero()
local mask = alpha:clone()
mask[torch.gt(mask, 0.0)] = 1
mask[torch.eq(mask, 0.0)] = 0
@ -26,10 +26,10 @@ function alpha_util.make_border(rgb, alpha, offset)
rgb[3][mask_nega] = 0
for i = 1, offset do
local mask_weight = alpha_util.sum2d:forward(mask:cuda()):float()
local mask_weight = sum2d:forward(mask:cuda()):float()
local border = rgb:clone()
for j = 1, 3 do
border[j]:copy(alpha_util.sum2d:forward(rgb[j]:reshape(1, rgb:size(2), rgb:size(3)):cuda()))
border[j]:copy(sum2d:forward(rgb[j]:reshape(1, rgb:size(2), rgb:size(3)):cuda()))
border[j]:cdiv((mask_weight + eps))
rgb[j][mask_nega] = border[j][mask_nega]
end