(self, is_fused_checkpoint_a, is_fused_checkpoint_b,
use_gpu_checkpoint_a, use_gpu_checkpoint_b,
use_gpu_test_a, use_gpu_test_b, freeze_mode)
| 165 | self.assertAllClose(test_loss, ref_loss, rtol=1.e-3, atol=1.e-3) |
| 166 | |
| 167 | def _testCheckpoint(self, is_fused_checkpoint_a, is_fused_checkpoint_b, |
| 168 | use_gpu_checkpoint_a, use_gpu_checkpoint_b, |
| 169 | use_gpu_test_a, use_gpu_test_b, freeze_mode): |
| 170 | batch, height, width, input_channels = 2, 4, 5, 3 |
| 171 | shape = [batch, height, width, input_channels] |
| 172 | base_path = '%s_%s_%s_%s_%s_%s' % (is_fused_checkpoint_a, |
| 173 | is_fused_checkpoint_b, |
| 174 | use_gpu_checkpoint_a, |
| 175 | use_gpu_checkpoint_b, use_gpu_test_a, |
| 176 | use_gpu_test_b) |
| 177 | |
| 178 | checkpoint_path_a = os.path.join(self.get_temp_dir(), |
| 179 | 'checkpoint_a_%s' % base_path) |
| 180 | self._train( |
| 181 | checkpoint_path_a, |
| 182 | shape, |
| 183 | use_gpu_checkpoint_a, |
| 184 | is_fused_checkpoint_a, |
| 185 | restore=False, |
| 186 | freeze_mode=freeze_mode) |
| 187 | checkpoint_path_b = os.path.join(self.get_temp_dir(), |
| 188 | 'checkpoint_b_%s' % base_path) |
| 189 | self._train( |
| 190 | checkpoint_path_b, |
| 191 | shape, |
| 192 | use_gpu_checkpoint_b, |
| 193 | is_fused_checkpoint_b, |
| 194 | restore=False, |
| 195 | freeze_mode=freeze_mode) |
| 196 | |
| 197 | vars_fused = self._train( |
| 198 | checkpoint_path_a, |
| 199 | shape, |
| 200 | use_gpu_test_a, |
| 201 | True, |
| 202 | restore=True, |
| 203 | freeze_mode=freeze_mode) |
| 204 | vars_nonfused = self._train( |
| 205 | checkpoint_path_b, |
| 206 | shape, |
| 207 | use_gpu_test_b, |
| 208 | False, |
| 209 | restore=True, |
| 210 | freeze_mode=freeze_mode) |
| 211 | self.assertEqual(len(vars_fused), 5) |
| 212 | self.assertEqual(len(vars_nonfused), 5) |
| 213 | for var_fused, var_nonfused in zip(vars_fused, vars_nonfused): |
| 214 | self.assertAllClose(var_fused, var_nonfused, atol=1e-5) |
| 215 | |
| 216 | image_val = np.random.rand(batch, height, width, |
| 217 | input_channels).astype(np.float32) |
| 218 | loss_fused_val = self._infer(checkpoint_path_a, image_val, shape, |
| 219 | use_gpu_test_a, True) |
| 220 | loss_nonfused_val = self._infer(checkpoint_path_b, image_val, shape, |
| 221 | use_gpu_test_b, False) |
| 222 | self.assertAllClose(loss_fused_val, loss_nonfused_val, atol=1e-6, rtol=3e-4) |
| 223 | |
| 224 | def _testCheckpointCrossDevice(self, ckpt_a_fused, ckpt_a_use_gpu, |
no test coverage detected