def forward(self, x, y):
h = F.leaky_relu(self.Conv1(x))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = F.leaky_relu(self.Conv2(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = F.leaky_relu(self.Conv3(h))
h = F.leaky_relu(self.Conv4(h))
h = F.leaky_relu(self.Conv5(h))
h = F.leaky_relu(self.Conv6(h))
h = F.leaky_relu(self.Conv7(h))
h = F.leaky_relu(self.FC1(h))
h = F.leaky_relu(self.FC2(h))
h = self.Output(h)
loss = F.mean_squared_error(h, y)
chainer.report({'loss': loss}, self)
return loss
def forward(self, x, y):
y = chainermn.functions.bcast(self.comm, y, 0)
partions = cp.array_split(x, self.comm.size, -2)
# This part needs fixing. Probably all conditions are not checked
if self.comm.rank == 0:
x = partions[0]
elif self.comm.rank == 1:
x = partions[1]
elif self.comm.rank == 2:
x = partions[2]
elif self.comm.rank == 3:
x = partions[3]
else:
print("Rank does not exist")
h = FX.halo_exchange_3d(self.comm, x, k_size=3, index=1, pad=0)
h = F.leaky_relu(self.Conv1(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = FX.halo_exchange_3d(self.comm, h, k_size=3, index=2, pad=0)
h = F.leaky_relu(self.Conv2(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
hs = chainermnx.functions.spatialallgather(self.comm, h)
h = F.concat(hs, -2)
h = F.leaky_relu(self.Conv3(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = F.leaky_relu(self.Conv4(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = F.leaky_relu(self.Conv5(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = F.leaky_relu(self.Conv6(h))
h = F.average_pooling_3d(h, ksize=2, stride=2)
h = F.leaky_relu(self.Conv7(h))
h = F.leaky_relu(self.FC1(h))
h = F.leaky_relu(self.FC2(h))
h = self.Output(h)
loss = F.mean_squared_error(h, y)
chainer.report({'loss': loss}, self)
# print("Rank ", self.comm.rank, "Completed forward and y is ", y)
return loss
def test_average_pooling_3d_invalid(self):
(x, ksize) = self._get_data(2)
with self.assertRaises(ValueError):
functions.average_pooling_3d(x, ksize)
def test_average_pooling_3d(self):
(x, ksize) = self._get_data(3)
testing.assert_allclose(
functions.average_pooling_nd(x, ksize).data,
functions.average_pooling_3d(x, ksize).data)
def downscale2x3d(h):
return F.average_pooling_3d(h, 2, 2, 0)
def test_average_pooling_3d_invalid(self):
(x, ksize) = self._get_data(2)
with pytest.raises(ValueError):
functions.average_pooling_3d(x, ksize)
def test_average_pooling_3d(self):
(x, ksize) = self._get_data(3)
testing.assert_allclose(
functions.average_pooling_nd(x, ksize).data,
functions.average_pooling_3d(x, ksize).data)
