return get_relay_op(name)(data0, data1)
return _impl
+def _squeeze():
+ def _impl(inputs, input_types):
+ data = inputs[0]
+ if len(inputs) == 1:
+ axis = None
+ else:
+ axis = [int(inputs[1])]
+
+ return _op.transform.squeeze(data, axis)
+ return _impl
+
def _unsqueeze():
def _impl(inputs, input_types):
data = inputs[0]
return _op.nn.max_pool1d(data, pool_size, strides, padding, "NCW", ceil_mode)
return _impl
+def _maxpool_3d():
+ def _impl(inputs, input_types):
+ data = inputs[0]
+
+ pool_size = _infer_shape(inputs[1])
+ strides = _infer_shape(inputs[2])
+ padding = _infer_shape(inputs[3])
+ dilation = _infer_shape(inputs[4])
+ ceil_mode = int(inputs[5])
+ if dilation != (1, 1, 1):
+ msg = "MaxPool3d with dilation %s is not implemented" % (str(dilation), )
+ raise NotImplementedError(msg)
+
+ return _op.nn.max_pool3d(data,
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ ceil_mode=ceil_mode)
+ return _impl
+
def _hardtanh():
def _impl(inputs, input_types):
a = inputs[0]
return _impl
+def _avg_pool3d():
+ def _impl(inputs, input_types):
+ data = inputs[0]
+
+ pool_size = _infer_shape(inputs[1])
+ if inputs[2]:
+ strides = _infer_shape(inputs[2])
+ else:
+ strides = pool_size
+ padding = _infer_shape(inputs[3])
+
+ ceil_mode = int(inputs[4])
+ count_include_pad = int(inputs[5])
+
+ return _op.nn.avg_pool3d(data,
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad)
+ return _impl
def _dropout():
def _impl(inputs, input_types):
"aten::ones" : _ones(),
"aten::zeros" : _zeros(),
"aten::to" : _to(),
+ "aten::squeeze" : _squeeze(),
"aten::unsqueeze" : _unsqueeze(),
"aten::cat" : _concatenate(),
"aten::slice" : _slice(),
"aten::max_pool2d" : _maxpool_2d(),
"aten::max_pool2d_with_indices" : _maxpool_2d(),
"aten::max_pool1d" : _maxpool_1d(),
+ "aten::max_pool3d" : _maxpool_3d(),
"aten::hardtanh" : _hardtanh(),
"aten::hardtanh_" : _hardtanh(),
"aten::_convolution" : _convolution(),
"aten::log_softmax" : _log_softmax(),
"aten::sigmoid" : _sigmoid(),
"aten::avg_pool2d" : _avg_pool2d(),
+ "aten::avg_pool3d" : _avg_pool3d(),
"aten::dropout" : _dropout(),
"aten::dropout_" : _dropout(),
"aten::feature_dropout" : _dropout(),
input_data = torch.rand(input_shape).float()
verify_model(Unsqueeze1().float().eval(), input_data=input_data)
+def test_forward_squeeze():
+ torch.set_grad_enabled(False)
+ input_shape = [2, 1, 10, 1, 10]
+
+ class Squeeze1(Module):
+ def forward(self, *args):
+ return args[0].squeeze()
+
+ class Squeeze2(Module):
+ def forward(self, *args):
+ return args[0].squeeze(1)
+
+ input_data = torch.rand(input_shape).float()
+ verify_model(Squeeze1().float().eval(), input_data=input_data)
+ verify_model(Squeeze2().float().eval(), input_data=input_data)
+
def test_forward_concatenate():
torch.set_grad_enabled(False)
input_shape = [1, 3, 10, 10]
stride=2).eval(),
input_data)
+def test_forward_maxpool3d():
+ torch.set_grad_enabled(False)
+ input_shape = [1, 3, 10, 10, 10]
+ input_data = torch.rand(input_shape).float()
+
+ verify_model(torch.nn.MaxPool3d(kernel_size=[1, 1, 1]).eval(),
+ input_data)
+ verify_model(torch.nn.MaxPool3d(kernel_size=[10, 10, 10]).eval(),
+ input_data)
+ verify_model(torch.nn.MaxPool3d(kernel_size=[4, 4, 4],
+ padding=2,
+ stride=2).eval(),
+ input_data)
+
def test_forward_split():
torch.set_grad_enabled(False)
input_shape = [4, 10]
verify_model(torch.nn.AvgPool2d(kernel_size=[10, 10]).eval(), input_data=input_data)
verify_model(AvgPool2D2().float().eval(), input_data=input_data)
+def test_forward_avgpool3d():
+ torch.set_grad_enabled(False)
+ input_shape = [1, 3, 10, 10, 10]
+
+ class AvgPool3D1(Module):
+ def forward(self, *args):
+ return torch.nn.functional.avg_pool3d(args[0], kernel_size=[10, 10, 10])
+
+ input_data = torch.rand(input_shape).float()
+ verify_model(torch.nn.AvgPool3d(kernel_size=[10, 10, 10]).eval(), input_data=input_data)
+ verify_model(AvgPool3D1().float().eval(), input_data=input_data)
+
def test_forward_hardtanh():
torch.set_grad_enabled(False)
input_shape = [10]
test_forward_add()
test_forward_subtract()
test_forward_multiply()
+ test_forward_squeeze()
test_forward_unsqueeze()
test_forward_concatenate()
test_forward_relu()
test_forward_adaptiveavgpool()
test_forward_maxpool2d()
test_forward_maxpool1d()
+ test_forward_maxpool3d()
test_forward_hardtanh()
test_forward_conv()
test_forward_conv_transpose()
test_forward_sigmoid()
test_forward_dense()
test_forward_avgpool()
+ test_forward_avgpool3d()
test_forward_dropout()
test_forward_slice()
test_forward_mean()