njuvision / awnet Goto Github PK

    mask = torch.autograd.Variable(torch.ones(x.size())).cuda()
    mask = nn.functional.grid_sample(mask, vgrid)

   # if W==128:
        # np.save('mask.npy', mask.cpu().data.numpy())
        # np.save('warp.npy', output.cpu().data.numpy())

    mask[mask<0.9999] = 0
    mask[mask>0] = 1

    return output*mask,mask

import torch
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
torch.cuda.set_device(0)
criteria = torch.nn.L1Loss()


# import awnet_pwcnet #please adjust as per your code

net = awnet_pwcnet.PWCDCNet().cuda()
#net = torch.nn.Conv2d(6, 3, 1).cuda()
optimizer = torch.optim.Adam(net.parameters(), lr = 0.1)

inp = Variable(torch.rand(1,3,256,448).to(device),requires_grad = True)
snd = Variable(torch.rand(1,3,256,448).to(device),requires_grad = True)
gt = Variable(torch.rand(1,2,256,448).to(device),requires_grad = True)

x = torch.cat((inp,snd),1)

# for layer in net.children():
#     for param in layer.parameters():
#         print(param)

for i in range(1,5):
    optimizer.zero_grad()
    out = net(x)
    out = F.interpolate(out,scale_factor=4,mode='bilinear',align_corners=False)*20
    loss = criteria(out,gt)
    print(out.shape)
    loss.backward()#retain_graph= True)
    print(net.state_dict())
    optimizer.step()