提交 e8494f72 编写于 作者: M Megvii Engine Team

fix(unittest): remove the torch dependency in test_conv

GitOrigin-RevId: 36dcab2b95d5bd596e7c3e8dc47ef658ad5f3b57
上级 779ed45d
......@@ -9,12 +9,9 @@
import itertools
import numpy as np
import pytest
import torch
import megengine as mge
from megengine import Parameter, tensor
from megengine.module import Conv2d, ConvTranspose2d
from megengine.module import ConvTranspose2d
from megengine.test import assertTensorClose
......@@ -37,28 +34,19 @@ def test_conv_transpose2d():
weight = np.random.normal(size=(IC, OC, KH, KW)).astype(np.float32)
bias = np.random.normal(size=(1, OC, 1, 1)).astype(np.float32)
# naive calculation use numpy
for n, ic, ih, iw in itertools.product(*map(range, [N, IC, IH, IW])):
oh, ow = ih * SH, iw * SW
out[n, :, oh : oh + KH, ow : ow + KW] += inp[n, ic, ih, iw] * weight[ic]
out = out[:, :, PH : OH - PH, PW : OW - PW]
if BIAS:
out += bias
# megengine conv_transpose2d calculation
conv_transpose2d = ConvTranspose2d(IC, OC, (KH, KW), (SH, SW), (PH, PW), bias=BIAS)
conv_transpose2d.weight = Parameter(weight, dtype=np.float32)
if BIAS:
conv_transpose2d.bias = Parameter(bias, dtype=np.float32)
y = conv_transpose2d(tensor(inp))
assertTensorClose(out, y.numpy(), max_err=2e-6)
torch_conv_transpose2d = torch.nn.ConvTranspose2d(
IC, OC, (KH, KW), stride=(SH, SW), padding=(PH, PW), bias=BIAS
)
torch_conv_transpose2d.weight = torch.nn.parameter.Parameter(torch.Tensor(weight))
if BIAS:
torch_conv_transpose2d.bias = torch.nn.parameter.Parameter(
torch.Tensor(bias).reshape(OC)
)
torch_y = torch_conv_transpose2d(torch.Tensor(inp))
assertTensorClose(torch_y.detach().numpy(), y.numpy(), max_err=2e-6)
assertTensorClose(out, y.numpy(), max_err=2e-6)
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册