|
- import zipfile
- import os
- import moxing as mox
- import json
- import time
-
- from mindspore import context
- from mindspore.context import ParallelMode
- from mindspore.communication.management import init, get_rank
-
-
- ### Copy multiple datasets from obs to training image ###
- def MultiObsToEnv(multi_data_url, data_dir):
- #--multi_data_url is json data, need to do json parsing for multi_data_url
- multi_data_json = json.loads(multi_data_url)
- for i in range(len(multi_data_json)):
- path = data_dir + "/" + multi_data_json[i]["dataset_name"]
- if not os.path.exists(path):
- os.makedirs(path)
- try:
- mox.file.copy_parallel(multi_data_json[i]["dataset_url"], path)
- print("Successfully Download {} to {}".format(multi_data_json[i]["dataset_url"],path))
- except Exception as e:
- print('moxing download {} to {} failed: '.format(
- multi_data_json[i]["dataset_url"], path) + str(e))
- #Set a cache file to determine whether the data has been copied to obs.
- #If this file exists during multi-card training, there is no need to copy the dataset multiple times.
- f = open("/cache/download_input.txt", 'w')
- f.close()
- try:
- if os.path.exists("/cache/download_input.txt"):
- print("download_input succeed")
- except Exception as e:
- print("download_input failed")
- return
- ### Copy the output model to obs ###
- def EnvToObs(train_dir, obs_train_url):
- try:
- mox.file.copy_parallel(train_dir, obs_train_url)
- print("Successfully Upload {} to {}".format(train_dir,
- obs_train_url))
- except Exception as e:
- print('moxing upload {} to {} failed: '.format(train_dir,
- obs_train_url) + str(e))
- return
- def DownloadFromQizhi(multi_data_url, data_dir):
- device_num = int(os.getenv('RANK_SIZE'))
- if device_num == 1:
- MultiObsToEnv(multi_data_url,data_dir)
- context.set_context(mode=context.GRAPH_MODE,device_target=args.device_target)
- if device_num > 1:
- # set device_id and init for multi-card training
- context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=int(os.getenv('ASCEND_DEVICE_ID')))
- context.reset_auto_parallel_context()
- context.set_auto_parallel_context(device_num = device_num, parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True, parameter_broadcast=True)
- init()
- #Copying obs data does not need to be executed multiple times, just let the 0th card copy the data
- local_rank=int(os.getenv('RANK_ID'))
- if local_rank%8==0:
- MultiObsToEnv(multi_data_url,data_dir)
- #If the cache file does not exist, it means that the copy data has not been completed,
- #and Wait for 0th card to finish copying data
- while not os.path.exists("/cache/download_input.txt"):
- time.sleep(1)
- return
- def UploadToQizhi(train_dir, obs_train_url):
- device_num = int(os.getenv('RANK_SIZE'))
- local_rank=int(os.getenv('RANK_ID'))
- if device_num == 1:
- EnvToObs(train_dir, obs_train_url)
- if device_num > 1:
- if local_rank%8==0:
- EnvToObs(train_dir, obs_train_url)
- return
-
-
- class LossFunc(nn.Cell):
- """loss function"""
- def __init__(self, attention_loss_weight=1.0, state='tuning'):
- super(LossFunc, self).__init__()
- self.arcface = ArcFace(81313)
- self._loss_fn = nn.SoftmaxCrossEntropyWithLogits(
- sparse=True, reduction="mean")
- self._autoencoder_loss_fn = nn.MSELoss()
- self.state = state
- self.attention_loss_weight = attention_loss_weight
-
- def construct(self, base, label):
- """construct"""
- label = ops.clip_by_value(label, 0, 81313)
- if self.state == 'tuning':
- base = self.arcface(base, label)
- total_loss = self._loss_fn(base, label)
- else:
- (desc_prelogits, attn_logits, stop_block3, dim_expanded_features) = base
- desc_prelogits = self.arcface(desc_prelogits, label)
- global_loss = self._loss_fn(desc_prelogits, label)
- attn_loss = self._loss_fn(attn_logits, label)
- autoencoder_loss = self._autoencoder_loss_fn(stop_block3, dim_expanded_features) * 10.0
- total_loss = self.attention_loss_weight * attn_loss + autoencoder_loss + global_loss
-
- return total_loss
-
-
- parser = argparse.ArgumentParser(description='MindSpore delf eval')
-
- parser.add_argument('--data_url',
- help='path to training/inference dataset folder',
- default= '/cache/data1/')
-
- parser.add_argument('--multi_data_url',
- help='path to multi dataset',
- default= '/cache/data/')
-
- parser.add_argument('--train_url',
- help='model folder to save/load',
- default= '/cache/output/')
-
- parser.add_argument('--device_target',
- type=str,
- default="Ascend",
- choices=['Ascend', 'CPU'],
- help='device where the code will be implemented (default: Ascend),if to use the CPU on the Qizhi platform:device_target=CPU')
-
- FLAGS = parser.parse_known_args()[0]
-
- if __name__ == '__main__':
- data_dir = '/cache/data'
- train_dir = '/cache/output'
- if not os.path.exists(data_dir):
- os.makedirs(data_dir)
- if not os.path.exists(train_dir):
- os.makedirs(train_dir)
-
- DownloadFromQizhi(FLAGS.multi_data_url, data_dir)
- dirpath = '/cache/data/GLDv2_Clean_MindRecord/MindRecord/'
- outFullName = '/cache/output/GLDv2_Clean_MindRecord.zip'
- zip = zipfile.ZipFile(outFullName, "w", zipfile.ZIP_DEFLATED)
- for path, dirnames, filenames in os.walk(dirpath):
- fpath = path.replace(dirpath, '')
- for filename in filenames:
- zip.write(os.path.join(path, filename), os.path.join(fpath, filename))
- zip.close()
- UploadToQizhi(train_dir,FLAGS.train_url)
|