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)