Browse Source

fix(unittest): remove the torch dependency in test_conv

GitOrigin-RevId: 36dcab2b95
tags/v0.3.2
Megvii Engine Team 5 years ago
parent
commit
e8494f729b
1 changed files with 5 additions and 17 deletions
  1. +5
    -17
      python_module/test/unit/module/test_conv.py

+ 5
- 17
python_module/test/unit/module/test_conv.py View File

@@ -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)

Loading…
Cancel
Save