def construct(self, tensor_a, tensor_b): a = F.mixed_precision_cast(mstype.float16, tensor_a) b = F.mixed_precision_cast(mstype.float16, tensor_b) c = self.sub(a, b) dictionary = {"key": a} result = self.net(c, key1=a, key2=dictionary) return result
def construct(self, image, target, weight, scale=None, center=None, score=None, idx=None): out = self._backbone(image) output = F.mixed_precision_cast(mstype.float32, out) target = F.mixed_precision_cast(mstype.float32, target) weight = F.mixed_precision_cast(mstype.float32, weight) return self._loss_fn(output, target, weight)
def construct(self, positions, forces, energy): outputs = self._network(positions) foutputs = -1 * self.grad_op(self._network)(positions) if self.add_cast_fp32: forces = F.mixed_precision_cast(ms.float32, forces) energy = F.mixed_precision_cast(ms.float32, energy) outputs = F.cast(outputs, ms.float32) if self._energy_fn is None: eloss = 0 else: eloss = self._energy_fn(outputs, energy) if self._force_fn is None: floss = 0 else: floss = self._force_fn(foutputs, forces) return eloss, floss, outputs, energy, foutputs, forces
def test_mixed_precision_cast(): x = Tensor(np.ones([2, 3], dtype=np.float32)) z = F.mixed_precision_cast(mstype.float16, x) assert z.dtype == mstype.float16
def construct(self, tensor_c, **kwargs): d = F.mixed_precision_cast(mstype.float16, tensor_c) dict_cast = F.mixed_precision_cast(mstype.float16, kwargs) e = self.add(d, dict_cast["key1"]) f = self.add(e, dict_cast["key2"]["key"]) return f
def construct(self, data, label): out = self._backbone(data) label = F.mixed_precision_cast(mstype.float32, label) return self._loss_fn(F.mixed_precision_cast(mstype.float32, out), label)