| 343 | return ret, grad, grad_estimate |
| 344 | |
| 345 | def _assertInferTensorChecks(self, op, grad_ops): |
| 346 | tmp_net = caffe2_pb2.NetDef() |
| 347 | tmp_net.op.extend([op]) |
| 348 | tmp_net.op.extend(grad_ops) |
| 349 | inferred_shapes, inferred_types = workspace.InferShapesAndTypes( |
| 350 | [tmp_net], |
| 351 | nets_proto=True, |
| 352 | ) |
| 353 | |
| 354 | outputs = set() |
| 355 | for grad_op in grad_ops: |
| 356 | outputs.update(grad_op.output) |
| 357 | |
| 358 | for output in outputs: |
| 359 | if output not in inferred_shapes: |
| 360 | raise Exception( |
| 361 | "expected output {} to be inferred".format(output)) |
| 362 | blob = workspace.FetchBlob(output) |
| 363 | correct_shape = list(blob.shape) |
| 364 | inferred_shape = list(inferred_shapes[output]) |
| 365 | if correct_shape != inferred_shape: |
| 366 | raise Exception( |
| 367 | "Mismatched inferred shape: want({}), got({})".format( |
| 368 | correct_shape, inferred_shape)) |
| 369 | |
| 370 | if type(blob) is np.ndarray: |
| 371 | if blob.dtype == np.dtype('float64'): |
| 372 | correct_type = caffe2_pb2.TensorProto.DOUBLE |
| 373 | elif blob.dtype == np.dtype('float32'): |
| 374 | correct_type = caffe2_pb2.TensorProto.FLOAT |
| 375 | elif blob.dtype == np.dtype('int32'): |
| 376 | correct_type = caffe2_pb2.TensorProto.INT32 |
| 377 | elif blob.dtype == np.dtype('int64'): |
| 378 | correct_type = caffe2_pb2.TensorProto.INT64 |
| 379 | else: |
| 380 | correct_type = "unknown {}".format(np.dtype) |
| 381 | else: |
| 382 | correct_type = str(type(blob)) |
| 383 | inferred_type = inferred_types[output] |
| 384 | if correct_type != inferred_type: |
| 385 | raise Exception( |
| 386 | "Mismatched inferred type: want({}), got({})".format( |
| 387 | correct_type, inferred_type)) |