| 179 | self._input_device_options = input_device_options |
| 180 | |
| 181 | def GetLossAndGrad( |
| 182 | self, op, grad_ops, inputs, input_names, input_to_check, grad_name, |
| 183 | outputs_with_grads |
| 184 | ): |
| 185 | for i in range(len(inputs)): |
| 186 | workspace.FeedBlob(input_names[i], inputs[i], |
| 187 | self._input_device_options.get( |
| 188 | input_names[i], self._device_option)) |
| 189 | x = inputs[input_to_check] |
| 190 | # Run. |
| 191 | workspace.RunOperatorOnce(op) |
| 192 | loss = 0. |
| 193 | # Get Loss and feed in the gradients, run gradient ops. |
| 194 | for idx in outputs_with_grads: |
| 195 | name = op.output[idx] |
| 196 | arr = workspace.FetchBlob(name) |
| 197 | loss += (arr**2).sum() |
| 198 | workspace.FeedBlob(name + '_grad', arr, self._device_option) |
| 199 | loss /= 2. |
| 200 | # Run gradient ops |
| 201 | workspace.RunOperatorsOnce(grad_ops) |
| 202 | # Get gradients |
| 203 | if isinstance(grad_name, core.GradientSlice): |
| 204 | workspace.FeedBlob('zeros', np.zeros_like(x, dtype=np.float32)) |
| 205 | workspace.FeedBlob('ones', np.ones(1, dtype=np.float32)) |
| 206 | gv_cpu_op = core.CreateOperator( |
| 207 | 'EnsureCPUOutput', grad_name.values, grad_name.values + '_cpu', |
| 208 | device_option=self._device_option |
| 209 | ) |
| 210 | gi_cpu_op = core.CreateOperator( |
| 211 | 'EnsureCPUOutput', grad_name.indices, grad_name.indices + '_cpu', |
| 212 | device_option=self._device_option |
| 213 | ) |
| 214 | sparse_to_dense_op = core.CreateOperator( |
| 215 | 'ScatterWeightedSum', |
| 216 | [ |
| 217 | 'zeros', 'ones', grad_name.indices + '_cpu', |
| 218 | grad_name.values + '_cpu', 'ones' |
| 219 | ], |
| 220 | 'zeros', |
| 221 | ) |
| 222 | workspace.RunOperatorOnce(gv_cpu_op) |
| 223 | workspace.RunOperatorOnce(gi_cpu_op) |
| 224 | workspace.RunOperatorOnce(sparse_to_dense_op) |
| 225 | grad = workspace.FetchBlob('zeros') |
| 226 | else: |
| 227 | grad = workspace.FetchBlob(grad_name) |
| 228 | return loss, grad |
| 229 | |
| 230 | def CheckSimple( |
| 231 | self, |