def do_nhwc_conv(test_case, with_cuda, with_bias):
x = flow.randn(2, 3, 4, 5)
conv = flow.nn.Conv2d(3, 4, 2, 1, bias=with_bias)
if with_cuda:
x = x.cuda()
conv.to("cuda")
eager_conv_x = flow.relu(conv(x)) + flow.relu(conv(x))
class GraphToRun(flow.nn.Graph):
def __init__(self):
super().__init__()
self.conv = conv
def build(self, x):
return flow.relu(self.conv(x)) + flow.relu(self.conv(x))
graph_to_run = GraphToRun()
lazy_conv_x = graph_to_run(x)
print(eager_conv_x.numpy().flatten()[:10])
print(lazy_conv_x.numpy().flatten()[:10])
test_case.assertTrue(
np.allclose(eager_conv_x.numpy(),
lazy_conv_x.numpy(),
rtol=1e-5,
atol=1e-5))
def do_nhwc_bacth_norm(test_case, with_cuda):
x = flow.randn(2, 3, 4, 5)
bn = flow.nn.BatchNorm2d(3)
if with_cuda:
x = x.cuda()
bn.to("cuda")
eager_batch_norm_res = flow.relu(bn(x))
class GraphToRun(flow.nn.Graph):
def __init__(self):
super().__init__()
self.m = bn
def build(self, x):
return flow.relu(self.m(x))
graph_to_run = GraphToRun()
lazy_batch_norm_res = graph_to_run(x)
test_case.assertTrue(
np.allclose(
eager_batch_norm_res.numpy(),
lazy_batch_norm_res.numpy(),
rtol=1e-5,
atol=1e-5,
))
def _relu_inplace(self):
return flow.relu(self, inplace=True)
def _relu(self, inplace=False):
return flow.relu(self, inplace=inplace)
def build(self, x):
return flow.relu(self.m(x))
def build(self, x):
return flow.relu(self.conv(x)) + flow.relu(self.conv(x))