def run(self): self.device.use() self.setup() while True: job, data = self.pipe.recv() if job == 'finalize': self.device.device.synchronize() break if job == 'update': # For reducing memory self.model.cleargrads() model = self.model model.cleargrads() x = self.converter(self.iterator.next(), self.device_index) batch_size = len(x) images = x['image'] viewpoints = x['viewpoint'] xp = model.xp representation, query_images, query_viewpoints = encode_scene( images, viewpoints, model, self.device_index) with self.reporter.scope({}): # pass dummy observation # Compute distribution parameterws (z_t_param_array, pixel_mean) = model.sample_z_and_x_params_from_posterior( query_images, query_viewpoints, representation) # Compute ELBO (ELBO, bits_per_pixel, negative_log_likelihood, kl_divergence) = estimate_ELBO(xp, query_images, z_t_param_array, pixel_mean, self.pixel_log_sigma, batch_size) # Update parameters loss = -ELBO loss.backward() del loss gg = gather_grads(self.model) nccl_data_type = _get_nccl_data_type(gg.dtype) null_stream = cuda.Stream.null self.comm.reduce(gg.data.ptr, gg.data.ptr, gg.size, nccl_data_type, nccl.NCCL_SUM, 0, null_stream.ptr) del gg self.model.cleargrads() gp = gather_params(self.model) nccl_data_type = _get_nccl_data_type(gp.dtype) self.comm.bcast(gp.data.ptr, gp.size, nccl_data_type, 0, null_stream.ptr) scatter_params(self.model, gp) del gp
def test_gather_scatter_params(self): cupy = cuda.cupy model0 = SimpleNet(dtype=self.dtype) model1 = SimpleNet(dtype=self.dtype) model0.to_gpu() model1.to_gpu() gp0 = mpu.gather_params(model0) mpu.scatter_params(model1, gp0) cupy.testing.assert_array_equal(model0.conv.W.data, model1.conv.W.data) cupy.testing.assert_array_equal(model0.conv.b.data, model1.conv.b.data) cupy.testing.assert_array_equal(model0.fc.W.data, model1.fc.W.data) cupy.testing.assert_array_equal(model0.fc.b.data, model1.fc.b.data)
def test_gather_scatter_params(self): cupy = cuda.cupy model0 = SimpleNet(dtype=self.dtype) model1 = SimpleNet(dtype=self.dtype) model0.to_gpu() model1.to_gpu() gp0 = mpu.gather_params(model0) mpu.scatter_params(model1, gp0) cupy.testing.assert_array_equal(model0.conv.W.data, model1.conv.W.data) cupy.testing.assert_array_equal(model0.conv.b.data, model1.conv.b.data) cupy.testing.assert_array_equal(model0.fc.W.data, model1.fc.W.data) cupy.testing.assert_array_equal(model0.fc.b.data, model1.fc.b.data)
def run(self): from cupy.cuda import nccl dev = cuda.Device(self.device) dev.use() self.setup() gp = None while True: job, data = self.pipe.recv() if job == 'finalize': dev.synchronize() break if job == 'update': # For reducing memory self.model.cleargrads() batch = self.iterator.next() x = converter_kaldi(batch[0], self.reader) observation = {} with self.reporter.scope(observation): loss = self.model(x) self.model.cleargrads() loss.backward() loss.unchain_backward() del loss gg = gather_grads(self.model) null_stream = cuda.Stream.null self.comm.reduce(gg.data.ptr, gg.data.ptr, gg.size, nccl.NCCL_FLOAT, nccl.NCCL_SUM, 0, null_stream.ptr) del gg self.model.cleargrads() gp = gather_params(self.model) self.comm.bcast(gp.data.ptr, gp.size, nccl.NCCL_FLOAT, 0, null_stream.ptr) scatter_params(self.model, gp) gp = None delete_feat(x)