mirror of https://github.com/Jittor/Jittor
48 lines
1.9 KiB
Python
48 lines
1.9 KiB
Python
# ***************************************************************
|
|
# Copyright (c) 2021 Jittor. All Rights Reserved.
|
|
# Maintainers:
|
|
# Haoyang Peng <2247838039@qq.com>
|
|
# Guoye Yang <498731903@qq.com>
|
|
# Dun Liang <randonlang@gmail.com>.
|
|
#
|
|
# This file is subject to the terms and conditions defined in
|
|
# file 'LICENSE.txt', which is part of this source code package.
|
|
# ***************************************************************
|
|
import unittest
|
|
import jittor as jt
|
|
import numpy as np
|
|
|
|
skip_this_test = False
|
|
try:
|
|
jt.dirty_fix_pytorch_runtime_error()
|
|
import torch
|
|
from torch.nn import functional as F
|
|
except:
|
|
torch = None
|
|
skip_this_test = True
|
|
|
|
|
|
@unittest.skipIf(skip_this_test, "No Torch Found")
|
|
class TestFoldOp(unittest.TestCase):
|
|
def test_fold(self):
|
|
# test unfold first and the test fold.
|
|
for i in range(4,10):
|
|
tn = np.random.randn(1,3,i,i).astype('float32')
|
|
ja = jt.array(tn)
|
|
ta = torch.autograd.Variable(torch.from_numpy(tn),requires_grad=True)
|
|
juf = jt.nn.unfold(ja,kernel_size=2,stride=2,dilation=2,padding=2)
|
|
tuf = F.unfold(ta,kernel_size=2,stride=2,dilation=2,padding=2)
|
|
assert np.allclose(juf.data,tuf.detach().numpy())
|
|
gjuf = jt.grad(juf,ja)
|
|
gtuf = torch.autograd.grad(tuf,ta,torch.ones_like(tuf),retain_graph=True)[0]
|
|
assert np.allclose(gjuf.data,gtuf.detach().numpy())
|
|
# test fold
|
|
jf = jt.nn.fold(juf,output_size=(i,i),kernel_size=2,stride=2,dilation=2,padding=2)
|
|
tf = F.fold(tuf,output_size=(i,i),kernel_size=2,stride=2,dilation=2,padding=2)
|
|
assert np.allclose(jf.data,tf.detach().numpy())
|
|
gjf = jt.grad(jf,juf)
|
|
gtf = torch.autograd.grad(tf,tuf,torch.ones_like(tf),retain_graph=True)[0]
|
|
assert np.allclose(gjf.data,gtf.detach().numpy())
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |