update

xt/model/ms_dist.py

@@ -137,6 +137,9 @@ def sample_dtype(self):
 
  def log_prob(self, x, logits):
  x = self.oneHot(x, self.size, self.on_value, self.off_value)
+ if x.dtype != logits.dtype:
+ logits = logits.astype(ms.float16)
+ x = x.astype(ms.float16)
  loss, _ = self.softmax_cross(logits, x)
  return -self.expand_dims(loss, -1)
 

xt/model/muzero/muzero_model_ms.py

@@ -36,8 +36,8 @@
 from zeus.common.util.register import Registers
 from mindspore import set_context
 from xt.model.dqn.dqn_cnn_ms import MyTrainOneStepCell
-set_context(runtime_num_threads=3, mode=0)
-
+set_context(runtime_num_threads=3, mode=0,device_id=0)
+from mindspore import amp
 # pylint: disable=W0201
 
 @Registers.model
@@ -105,6 +105,7 @@ def __init__(self, model_info):
  device_target = ms.get_context("device_target")
  if device_target == 'Ascend':
  manager = FixedLossScaleUpdateCell(loss_scale_value=2**14)
+ self.net_with_loss = amp.auto_mixed_precision(self.net_with_loss, "O2")
  self.train_net = MyTrainOneStepCell(self.net_with_loss, self.adam, manager)
  elif device_target == "GPU" or device_target == "CPU" :
  self.train_net = myTrainOneStepCell(self.net_with_loss, optimizer=self.adam)

xt/model/ppo/ppo_ms.py

@@ -32,8 +32,8 @@
 from xt.model.dqn.dqn_cnn_ms import MyTrainOneStepCell
 from mindspore.train import Model
 
-
-ms.set_context(runtime_num_threads=5,mode =0)
+from mindspore import amp
+ms.set_context(runtime_num_threads=5,mode =0,device_id=0)
 
 @Registers.model
 class PPOMS(XTModel_MS):
@@ -79,6 +79,7 @@ def __init__(self, model_info):
  device_target = ms.get_context("device_target")
  if device_target == 'Ascend':
  manager = FixedLossScaleUpdateCell(loss_scale_value=2**14)
+ forward_fn = amp.auto_mixed_precision(forward_fn, "O2")
  self.train_net = MyTrainOneStepCell(forward_fn, adam, manager, grad_clip=True, clipnorm=self._max_grad_norm)
  elif device_target == "GPU" or device_target == "CPU":
  self.train_net = myTrainOneStepCell(forward_fn, optimizer=adam, max_grad_norm=self._max_grad_norm)
@@ -143,6 +144,7 @@ def __init__(self, net, loss_fn, dist):
  def construct(self, state_ph, adv_ph, old_logp_ph, behavior_action, target_v, old_v_ph):
  pi_latent, v_out = self.net(state_ph)
  ent = self.dist.entropy(pi_latent)
+ pi_latent = pi_latent.astype(ms.float16)
  action_log_prob = self.dist.log_prob(behavior_action, pi_latent)
  loss = self._loss_fn(action_log_prob, ent, adv_ph, old_logp_ph, target_v, v_out, old_v_ph)
  return loss