def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size):
"""PS embeddingLookup cache enable set."""
self.embedding_table.cache_enable = True
self.embedding_table.is_param_ps = True
_set_cache_enable(True)
if _is_role_worker():
_insert_hash_table_size(self.embedding_table.name, vocab_cache_size, embedding_size, vocab_size)
def _process_vocab_cache(self, slice_mode):
"""PS embeddingLookup cache check and process."""
self.cache_enable = False
if self.vocab_cache_size > 0:
if self.target == 'CPU':
logger.warning("The configuration of 'vocab_cache_size' is valid only in 'DEVICE' target, "
"current target is CPU, so it will be ignored.")
return
enable_ps = _get_ps_context("enable_ps")
if not enable_ps:
logger.warning("The configuration of 'vocab_cache_size' is valid only in parameter server trainning "
"mode, current mode is not parameter server trainning mode, so it will be ignored.")
return
parallel_mode = _get_parallel_mode()
is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
if is_auto_parallel:
device_num = get_group_size()
full_batch = _get_full_batch()
if device_num > 1 and not (full_batch and slice_mode == "table_row_slice"):
raise ValueError("The embeddingLookup cache of parameter server parallel only be used "
"in 'full_batch' and 'table_row_slice' parallel strategy.")
self.vocab_cache_size = self.vocab_cache_size * device_num
self.cache_enable = True
if _is_role_worker():
self.vocab_size = self.vocab_cache_size
def _process_vocab_cache(self, slice_mode):
"""PS embeddingLookup cache check and process."""
self.cache_enable = False
if self.vocab_cache_size > 0:
if self.target == 'CPU':
logger.warning("The configuration of 'vocab_cache_size' is valid only in 'DEVICE' target, "
"current target is CPU, so it will be ignored.")
return
enable_ps = _get_ps_context("enable_ps")
if not enable_ps:
logger.warning("The configuration of 'vocab_cache_size' is valid only in parameter server trainning "
"mode, current mode is not parameter server trainning mode, so it will be ignored.")
return
parallel_mode = _get_parallel_mode()
is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
if is_auto_parallel:
rank_size = get_group_size()
rank_id = get_rank()
full_batch = _get_full_batch()
if rank_size > 1 and not (full_batch and slice_mode == "table_row_slice"):
raise ValueError("The embeddingLookup cache of parameter server parallel only be used "
"in 'full_batch' and 'table_row_slice' parallel strategy.")
self.vocab_cache_size = self.vocab_cache_size * rank_size
_set_rank_id(rank_id)
self.cache_enable = True
if _is_role_worker():
self.vocab_size = self.vocab_cache_size
if context.get_context("enable_sparse") != self.sparse:
raise ValueError("The value of parameter 'sparse' must be same for all EmbeddingLookup "
"kernels and equal the value of 'enable_sparse' in context setting in "
"parameter server cache mode")
def train_and_eval(config):
"""
test_train_eval
"""
set_seed(1000)
data_path = config.data_path
batch_size = config.batch_size
epochs = config.epochs
if config.dataset_type == "tfrecord":
dataset_type = DataType.TFRECORD
elif config.dataset_type == "mindrecord":
dataset_type = DataType.MINDRECORD
else:
dataset_type = DataType.H5
parameter_server = bool(config.parameter_server)
if cache_enable:
config.full_batch = True
print("epochs is {}".format(epochs))
if config.full_batch:
context.set_auto_parallel_context(full_batch=True)
ds.config.set_seed(1)
ds_train = create_dataset(data_path,
train_mode=True,
epochs=1,
batch_size=batch_size * get_group_size(),
data_type=dataset_type)
ds_eval = create_dataset(data_path,
train_mode=False,
epochs=1,
batch_size=batch_size * get_group_size(),
data_type=dataset_type)
else:
ds_train = create_dataset(data_path,
train_mode=True,
epochs=1,
batch_size=batch_size,
rank_id=get_rank(),
rank_size=get_group_size(),
data_type=dataset_type)
ds_eval = create_dataset(data_path,
train_mode=False,
epochs=1,
batch_size=batch_size,
rank_id=get_rank(),
rank_size=get_group_size(),
data_type=dataset_type)
print("ds_train.size: {}".format(ds_train.get_dataset_size()))
print("ds_eval.size: {}".format(ds_eval.get_dataset_size()))
net_builder = ModelBuilder()
train_net, eval_net = net_builder.get_net(config)
train_net.set_train()
auc_metric = AUCMetric()
model = Model(train_net,
eval_network=eval_net,
metrics={"auc": auc_metric})
if cache_enable:
config.stra_ckpt = os.path.join(
config.stra_ckpt + "-{}".format(get_rank()), "strategy.ckpt")
context.set_auto_parallel_context(
strategy_ckpt_save_file=config.stra_ckpt)
eval_callback = EvalCallBack(model, ds_eval, auc_metric, config)
callback = LossCallBack(config=config)
if _is_role_worker():
if cache_enable:
ckptconfig = CheckpointConfig(
save_checkpoint_steps=ds_train.get_dataset_size() * epochs,
keep_checkpoint_max=1,
integrated_save=False)
else:
ckptconfig = CheckpointConfig(
save_checkpoint_steps=ds_train.get_dataset_size(),
keep_checkpoint_max=5)
else:
ckptconfig = CheckpointConfig(save_checkpoint_steps=1,
keep_checkpoint_max=1)
ckpoint_cb = ModelCheckpoint(prefix='widedeep_train',
directory=config.ckpt_path + '/ckpt_' +
str(get_rank()) + '/',
config=ckptconfig)
callback_list = [
TimeMonitor(ds_train.get_dataset_size()), eval_callback, callback
]
if get_rank() == 0:
callback_list.append(ckpoint_cb)
model.train(epochs,
ds_train,
callbacks=callback_list,
dataset_sink_mode=bool(parameter_server and cache_enable))
def train_and_eval(config):
"""
test_train_eval
"""
set_seed(1000)
data_path = config.data_path
batch_size = config.batch_size
epochs = config.epochs
if config.dataset_type == "tfrecord":
dataset_type = DataType.TFRECORD
elif config.dataset_type == "mindrecord":
dataset_type = DataType.MINDRECORD
else:
dataset_type = DataType.H5
parameter_server = bool(config.parameter_server)
cache_enable = config.vocab_cache_size > 0
print("epochs is {}".format(epochs))
ds_train = create_dataset(data_path,
train_mode=True,
epochs=1,
batch_size=batch_size,
data_type=dataset_type)
ds_eval = create_dataset(data_path,
train_mode=False,
epochs=1,
batch_size=batch_size,
data_type=dataset_type)
print("ds_train.size: {}".format(ds_train.get_dataset_size()))
print("ds_eval.size: {}".format(ds_eval.get_dataset_size()))
net_builder = ModelBuilder()
train_net, eval_net = net_builder.get_net(config)
train_net.set_train()
auc_metric = AUCMetric()
model = Model(train_net,
eval_network=eval_net,
metrics={"auc": auc_metric})
eval_callback = EvalCallBack(model, ds_eval, auc_metric, config)
callback = LossCallBack(config=config)
if _is_role_worker():
if cache_enable:
ckptconfig = CheckpointConfig(
save_checkpoint_steps=ds_train.get_dataset_size() * epochs,
keep_checkpoint_max=1)
else:
ckptconfig = CheckpointConfig(
save_checkpoint_steps=ds_train.get_dataset_size(),
keep_checkpoint_max=5)
else:
ckptconfig = CheckpointConfig(save_checkpoint_steps=1,
keep_checkpoint_max=1)
ckpoint_cb = ModelCheckpoint(prefix='widedeep_train',
directory=config.ckpt_path,
config=ckptconfig)
callback_list = [
TimeMonitor(ds_train.get_dataset_size()), eval_callback, callback,
ckpoint_cb
]
model.train(epochs,
ds_train,
callbacks=callback_list,
dataset_sink_mode=(parameter_server and cache_enable))
def __init__(self,
vocab_size,
embedding_size,
param_init='normal',
target='CPU',
slice_mode='batch_slice',
manual_shapes=None,
max_norm=None,
sparse=True,
vocab_cache_size=0):
super(EmbeddingLookup, self).__init__()
validator.check_value_type('sparse', sparse, [bool], self.cls_name)
self.target = target
if target not in ('CPU', 'DEVICE'):
raise ValueError(
'Attr \'target\' of \'EmbeddingLookup\' Op passed ' +
str(target) +
', should be one of values in \'CPU\', \'DEVICE\'.')
if not sparse and target == 'CPU':
raise ValueError(
'When target is CPU, embedding_lookup must be sparse.')
enable_ps = context.get_ps_context("enable_ps")
if not enable_ps and vocab_cache_size > 0:
logger.warning(
"The configuration of 'vocab_cache_size' is valid only in parameter server trainning mode, "
"current mode is not parameter server trainning mode, so it will be ignored."
)
vocab_cache_size = 0
if sparse:
self.gatherv2 = P.SparseGatherV2()
else:
self.gatherv2 = P.GatherV2()
self.embeddinglookup = P.EmbeddingLookup().add_prim_attr(
'primitive_target', 'CPU')
self.vocab_size = validator.check_positive_int(vocab_size,
'vocab_size')
self.vocab_cache_size = validator.check_non_negative_int(
vocab_cache_size, 'vocab_cache_size')
self.embedding_size = validator.check_positive_int(
embedding_size, 'embedding_size')
parallel_mode = _get_parallel_mode()
is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL,
ParallelMode.AUTO_PARALLEL)
self.cache_enable = self.vocab_cache_size > 0
if self.cache_enable:
if is_auto_parallel:
self.vocab_cache_size = self.vocab_cache_size * get_group_size(
)
self.vocab_size = self.vocab_cache_size
self.embedding_table = Parameter(initializer(
param_init, [self.vocab_size, self.embedding_size]),
name='embedding_table')
if self.cache_enable:
self.embedding_table.cache_enable = True
_set_cache_enable(True)
if _is_role_worker():
_insert_hash_table_size(self.embedding_table.name,
vocab_cache_size, embedding_size,
vocab_size)
self.forward_unique = False
self.gather_revert = P.GatherV2()
self.unique = P.Unique().shard(((1, ), ))
self.reshape = P.Reshape()
self.shape = P.Shape()
indices_shape_size = 2
if slice_mode == "field_slice" and is_auto_parallel:
if not manual_shapes:
raise ValueError(
"in slice field mode, the manual_shapes should not be none"
)
if not isinstance(manual_shapes, tuple):
raise TypeError(
"manual_shapes type must be tuple(int) cannot be {}!".
format(type(manual_shapes)))
for dim in manual_shapes:
validator.check_positive_int(dim, 'manual shape dim',
self.cls_name)
self.gatherv2.add_prim_attr("manual_split", manual_shapes)
self.embeddinglookup.add_prim_attr("manual_split", manual_shapes)
self.gatherv2.shard(((get_group_size(), 1), (1, get_group_size())))
self.embeddinglookup.shard(
((get_group_size(), 1), (1, get_group_size())))
elif slice_mode == "table_row_slice" and is_auto_parallel:
if target == 'DEVICE' and not self.cache_enable:
indices_shape_size = 1
self.gather_revert.shard(((1, 1), (get_group_size(), )))
self.forward_unique = True
indices_strategy = (1, ) * indices_shape_size
self.gatherv2.shard(((get_group_size(), 1), indices_strategy))
self.embeddinglookup.shard(
((get_group_size(), 1), indices_strategy))
elif slice_mode == "table_column_slice" and is_auto_parallel:
if target == 'DEVICE':
indices_shape_size = 1
self.gather_revert.shard(((1, get_group_size()), (1, )))
self.forward_unique = True
indices_strategy = (1, ) * indices_shape_size
self.gatherv2.shard(((1, get_group_size()), indices_strategy))
self.embeddinglookup.shard(
((1, get_group_size()), indices_strategy))
elif slice_mode == "batch_slice" and is_auto_parallel:
indices_strategy = [get_group_size()]
indices_strategy.extend([1] * (indices_shape_size - 1))
indices_strategy = tuple(indices_strategy)
self.gatherv2.shard(((1, 1), indices_strategy))
self.embeddinglookup.shard(((1, 1), indices_strategy))
else:
if is_auto_parallel:
raise ValueError(
"slice_mode should support mode in nn.EmbeddingLookup, but get "
+ str(slice_mode))
self.embedding_table.unique = self.forward_unique
self.max_norm = max_norm
if self.max_norm is not None:
self.max_norm = validator.check_positive_float(
self.max_norm, 'max_norm', self.cls_name)
self.max_norm = Tensor(self.max_norm, dtype=mstype.float32)
criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
net_with_criterion = WithLossCell(net, criterion)
train_network = TrainOneStepCell(net_with_criterion, optimizer)
train_network.set_train()
losses = []
for _ in range(epoch):
data = Tensor(np.random.randint(0, 15, (32, 3), np.int32))
label = Tensor(np.random.randint(0, 9, (32), np.int32))
if _is_role_pserver():
train_network(data, label)
sys.exit()
else:
loss = train_network(data, label).asnumpy()
losses.append(loss)
print(losses)
return losses
envs = os.environ
if __name__ == "__main__":
set_seed(0)
ps_loss = do_sparse_embedding(True)
if _is_role_worker():
context.reset_ps_context()
set_seed(0)
no_ps_loss = do_sparse_embedding()
context.set_ps_context(enable_ps=True)
assert np.allclose(ps_loss, no_ps_loss, rtol=1.0e-6, atol=1.0e-6)