def __init__(self, strategy1):
super().__init__()
self.matmul1 = P.MatMul().set_strategy(strategy1)
self.weight = Parameter(Tensor(np.ones([512, 256]).astype(np.float32) * 0.01), "w", requires_grad=True)
self.matmul2 = P.MatMul()
def __init__(self):
super(Net, self).__init__()
self.weight = Parameter(Tensor(np.ones([64, 10]).astype(np.float32)), name="weight")
self.bias = Parameter(Tensor(np.ones([10]).astype(np.float32)), name="bias")
self.matmul = P.MatMul()
self.biasAdd = P.BiasAdd()
def create_predict_data():
"""user-defined predict data"""
inputs_np = np.random.randn(128, 96).astype(np.float32)
return Tensor(inputs_np)
def __init__(self, controller):
super(Dihedral, self).__init__()
self.CONSTANT_Pi = 3.1415926535897932
if controller.amber_parm is not None:
file_path = controller.amber_parm
self.read_information_from_amberfile(file_path)
self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32)
self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32)
self.atom_c = Tensor(np.asarray(self.h_atom_c, np.int32), mstype.int32)
self.atom_d = Tensor(np.asarray(self.h_atom_d, np.int32), mstype.int32)
self.pk = Tensor(np.asarray(self.pk, np.float32), mstype.float32)
self.gamc = Tensor(np.asarray(self.gamc, np.float32), mstype.float32)
self.gams = Tensor(np.asarray(self.gams, np.float32), mstype.float32)
self.pn = Tensor(np.asarray(self.pn, np.float32), mstype.float32)
self.ipn = Tensor(np.asarray(self.ipn, np.int32), mstype.int32)
self.mul_weight = Parameter(initializer("ones", (8, 8, 8), ms.float32),
name="mul_weight")
self.matmul_weight = Parameter(initializer("ones", (64, 16),
ms.float32),
name="matmul_weight")
self.axis = axis
def construct(self, x, b):
out = self.gatherv2(self.param, x, self.axis)
out = self.mul(out, self.mul_weight)
out = self.reshape(out, (8, 64))
out = self.matmul(out, self.matmul_weight)
return out
_x = Tensor(np.ones([8, 8]), dtype=ms.int32)
_b = Tensor(np.ones([64, 8]), dtype=ms.float32)
def compile_net(net):
context.set_context(save_graphs=True)
optimizer = Momentum(net.trainable_params(),
learning_rate=0.1,
momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, _x, _b, auto_parallel_mode=True)
context.reset_auto_parallel_context()
def __init__(self):
super(SoftmaxCrossEntropyWithLogitsNet, self).__init__()
self.soft = P.SoftmaxCrossEntropyWithLogits()
self.value = (Tensor(np.zeros((2, 2)).astype(np.float32)),
Tensor(np.ones((2, 2)).astype(np.float32)))
choices=["Ascend", "GPU", "CPU"],
default="Ascend",
help="device target")
args = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target=args.device_target,
device_id=args.device_id)
if __name__ == "__main__":
data_config = DataConfig()
model_builder = ModelBuilder(ModelConfig, TrainConfig)
_, network = model_builder.get_train_eval_net()
load_checkpoint(args.ckpt_file, net=network)
batch_ids = Tensor(
np.zeros([data_config.batch_size,
data_config.data_field_size]).astype(np.int32))
batch_wts = Tensor(
np.zeros([data_config.batch_size,
data_config.data_field_size]).astype(np.float32))
labels = Tensor(np.zeros([data_config.batch_size, 1]).astype(np.float32))
input_data = [batch_ids, batch_wts, labels]
export(network,
*input_data,
file_name=args.file_name,
file_format=args.file_format)
def __init__(self, max_cycles=10):
super(ForwardNet, self).__init__()
self.max_cycles = max_cycles
self.zero = Tensor(np.array(0), mstype.int32)
self.i = Tensor(np.array(0), mstype.int32)
self.weight = Parameter(Tensor(np.array(0), mstype.int32))
def test_net():
x = np.array([1.0, 4.0, 9.0]).astype(np.float32)
square = Net()
output = square(Tensor(x))
expect = np.array([1.0, 16.0, 81.0]).astype(np.float32)
assert (output.asnumpy() == expect).all()
def test_mul():
x0_np = np.random.uniform(-2, 2, (2, 3, 4, 4)).astype(np.float32)
y0_np = np.random.uniform(-2, 2, (2, 3, 4, 4)).astype(np.float32)
x1_np = np.random.uniform(-2, 2, (2, 3, 4, 4)).astype(np.float32)
y1_np = np.random.uniform(-2, 2, (2, 1, 4, 4)).astype(np.float32)
x2_np = np.random.uniform(-2, 2, (2, 1, 1, 4)).astype(np.float32)
y2_np = np.random.uniform(-2, 2, (2, 3, 4, 4)).astype(np.float32)
x3_np = np.random.uniform(-2, 2, 1).astype(np.float32)
y3_np = np.random.uniform(-2, 2, 1).astype(np.float32)
x4_np = np.array(768).astype(np.float32)
y4_np = np.array(3072.5).astype(np.float32)
x0 = Tensor(x0_np)
y0 = Tensor(y0_np)
x1 = Tensor(x1_np)
y1 = Tensor(y1_np)
x2 = Tensor(x2_np)
y2 = Tensor(y2_np)
x3 = Tensor(x3_np)
y3 = Tensor(y3_np)
x4 = Tensor(x4_np)
y4 = Tensor(y4_np)
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
mul = NetMul()
output0 = mul(x0, y0)
expect0 = np.multiply(x0_np, y0_np)
diff0 = output0.asnumpy() - expect0
error0 = np.ones(shape=expect0.shape) * 1.0e-5
assert np.all(diff0 < error0)
assert output0.shape() == expect0.shape
output1 = mul(x1, y1)
expect1 = np.multiply(x1_np, y1_np)
diff1 = output1.asnumpy() - expect1
error1 = np.ones(shape=expect1.shape) * 1.0e-5
assert np.all(diff1 < error1)
assert output1.shape() == expect1.shape
output2 = mul(x2, y2)
expect2 = np.multiply(x2_np, y2_np)
diff2 = output2.asnumpy() - expect2
error2 = np.ones(shape=expect2.shape) * 1.0e-5
assert np.all(diff2 < error2)
assert output2.shape() == expect2.shape
output3 = mul(x3, y3)
expect3 = np.multiply(x3_np, y3_np)
diff3 = output3.asnumpy() - expect3
error3 = np.ones(shape=expect3.shape) * 1.0e-5
assert np.all(diff3 < error3)
assert output3.shape() == expect3.shape
output4 = mul(x4, y4)
expect4 = np.multiply(x4_np, y4_np)
diff4 = output4.asnumpy() - expect4
error4 = np.ones(shape=expect4.shape) * 1.0e-5
assert np.all(diff4 < error4)
assert output4.shape() == expect4.shape
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
mul = NetMul()
output0 = mul(x0, y0)
expect0 = np.multiply(x0_np, y0_np)
diff0 = output0.asnumpy() - expect0
error0 = np.ones(shape=expect0.shape) * 1.0e-5
assert np.all(diff0 < error0)
assert output0.shape() == expect0.shape
output1 = mul(x1, y1)
expect1 = np.multiply(x1_np, y1_np)
diff1 = output1.asnumpy() - expect1
error1 = np.ones(shape=expect1.shape) * 1.0e-5
assert np.all(diff1 < error1)
assert output1.shape() == expect1.shape
output2 = mul(x2, y2)
expect2 = np.multiply(x2_np, y2_np)
diff2 = output2.asnumpy() - expect2
error2 = np.ones(shape=expect2.shape) * 1.0e-5
assert np.all(diff2 < error2)
assert output2.shape() == expect2.shape
output3 = mul(x3, y3)
expect3 = np.multiply(x3_np, y3_np)
diff3 = output3.asnumpy() - expect3
error3 = np.ones(shape=expect3.shape) * 1.0e-5
assert np.all(diff3 < error3)
assert output3.shape() == expect3.shape
output4 = mul(x4, y4)
expect4 = np.multiply(x4_np, y4_np)
diff4 = output4.asnumpy() - expect4
error4 = np.ones(shape=expect4.shape) * 1.0e-5
assert np.all(diff4 < error4)
assert output4.shape() == expect4.shape
def test_net():
x = np.random.random(size=(2, 3)).astype(np.float32)
sigmoid = Net()
output = sigmoid(Tensor(x))
print("=================output====================")
print(output)
device_target=args.device_target)
if args.preprocess == "true":
print("============== Starting Data Pre-processing ==============")
convert_to_mindrecord(cfg.embed_size, args.aclimdb_path,
args.preprocess_path, args.glove_path)
embedding_table = np.loadtxt(
os.path.join(args.preprocess_path, "weight.txt")).astype(np.float32)
network = SentimentNet(vocab_size=embedding_table.shape[0],
embed_size=cfg.embed_size,
num_hiddens=cfg.num_hiddens,
num_layers=cfg.num_layers,
bidirectional=cfg.bidirectional,
num_classes=cfg.num_classes,
weight=Tensor(embedding_table),
batch_size=cfg.batch_size)
# pre_trained
if args.pre_trained:
load_param_into_net(network, load_checkpoint(args.pre_trained))
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
opt = nn.Momentum(network.trainable_params(), cfg.learning_rate,
cfg.momentum)
loss_cb = LossMonitor()
model = Model(network, loss, opt, {'acc': Accuracy()})
print("============== Starting Training ==============")
ds_train = lstm_create_dataset(args.preprocess_path, cfg.batch_size, 1)
config_ck = CheckpointConfig(
def train_on_ascend():
config = config_ascend_quant if args_opt.quantization_aware else config_ascend
print("training args: {}".format(args_opt))
print("training configure: {}".format(config))
print("parallel args: rank_id {}, device_id {}, rank_size {}".format(rank_id, device_id, rank_size))
epoch_size = config.epoch_size
# distribute init
if run_distribute:
context.set_auto_parallel_context(device_num=rank_size,
parallel_mode=ParallelMode.DATA_PARALLEL,
parameter_broadcast=True,
mirror_mean=True)
init()
# define network
network = mobilenetV2(num_classes=config.num_classes)
# define loss
if config.label_smooth > 0:
loss = CrossEntropyWithLabelSmooth(smooth_factor=config.label_smooth, num_classes=config.num_classes)
else:
loss = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction='mean')
# define dataset
dataset = create_dataset(dataset_path=args_opt.dataset_path,
do_train=True,
config=config,
device_target=args_opt.device_target,
repeat_num=1,
batch_size=config.batch_size)
step_size = dataset.get_dataset_size()
# load pre trained ckpt
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
load_param_into_net(network, param_dict)
# convert fusion network to quantization aware network
if config.quantization_aware:
network = quant.convert_quant_network(network,
bn_fold=True,
per_channel=[True, False],
symmetric=[True, False])
# get learning rate
lr = Tensor(get_lr(global_step=config.start_epoch * step_size,
lr_init=0,
lr_end=0,
lr_max=config.lr,
warmup_epochs=config.warmup_epochs,
total_epochs=epoch_size + config.start_epoch,
steps_per_epoch=step_size))
# define optimization
opt = nn.Momentum(filter(lambda x: x.requires_grad, network.get_parameters()), lr, config.momentum,
config.weight_decay)
# define model
model = Model(network, loss_fn=loss, optimizer=opt)
print("============== Starting Training ==============")
callback = None
if rank_id == 0:
callback = [Monitor(lr_init=lr.asnumpy())]
if config.save_checkpoint:
config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_epochs * step_size,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="mobilenetV2",
directory=config.save_checkpoint_path,
config=config_ck)
callback += [ckpt_cb]
model.train(epoch_size, dataset, callbacks=callback)
print("============== End Training ==============")
params = tfm_model.trainable_params()
weights = load_infer_weights(config)
for param in params:
value = param.data
weights_name = param.name
if weights_name not in weights:
raise ValueError(f"{weights_name} is not found in weights.")
if isinstance(value, Tensor):
if weights_name in weights:
assert weights_name in weights
if isinstance(weights[weights_name], Parameter):
if param.data.dtype == "Float32":
param.set_data(
Tensor(weights[weights_name].data.asnumpy(),
mstype.float32))
elif param.data.dtype == "Float16":
param.set_data(
Tensor(weights[weights_name].data.asnumpy(),
mstype.float16))
elif isinstance(weights[weights_name], Tensor):
param.set_data(
Tensor(weights[weights_name].asnumpy(), config.dtype))
elif isinstance(weights[weights_name], np.ndarray):
param.set_data(Tensor(weights[weights_name], config.dtype))
else:
param.set_data(weights[weights_name])
else:
print("weight not found in checkpoint: " + weights_name)
param.set_data(zero_weight(value.asnumpy().shape))
def test_net():
features = np.random.randn(32, 1001).astype(np.float16)
labels = np.random.randn(32).astype(np.int32)
SparseSoftmaxCrossEntropyWithLogits = Net()
output = SparseSoftmaxCrossEntropyWithLogits(Tensor(features), Tensor(labels))
print(output.asnumpy())
def zeros(shape):
"""Create zeros like shape."""
return Tensor(np.zeros(tuple(shape), np.float32))
return task_params[self.task_name]["num_labels"]
return DEFAULT_NUM_LABELS
@property
def seq_length(self):
if self.task_name in task_params and "seq_length" in task_params[self.task_name]:
return task_params[self.task_name]["seq_length"]
return DEFAULT_SEQ_LENGTH
if __name__ == '__main__':
task = Task(args.task_name)
td_student_net_cfg.seq_length = task.seq_length
td_student_net_cfg.batch_size = DEFAULT_BS
eval_model = BertModelCLS(td_student_net_cfg, False, task.num_labels, 0.0, phase_type="student")
param_dict = load_checkpoint(args.ckpt_file)
new_param_dict = {}
for key, value in param_dict.items():
new_key = re.sub('tinybert_', 'bert_', key)
new_key = re.sub('^bert.', '', new_key)
new_param_dict[new_key] = value
load_param_into_net(eval_model, new_param_dict)
eval_model.set_train(False)
input_ids = Tensor(np.zeros((td_student_net_cfg.batch_size, task.seq_length), np.int32))
token_type_id = Tensor(np.zeros((td_student_net_cfg.batch_size, task.seq_length), np.int32))
input_mask = Tensor(np.zeros((td_student_net_cfg.batch_size, task.seq_length), np.int32))
export(eval_model, input_ids, token_type_id, input_mask, file_name=args.output_file, file_format="AIR")
def random_normal(*size):
"""Apply random values from a normal distribution."""
# return P.StandardNormal()(size)
return Tensor(np.random.randn(*size).astype(np.float32))
def __init__(self):
super(ConstDepend, self).__init__()
self.value = (Tensor(np.zeros((2, 3)).astype(np.float32)),
Tensor(np.ones((2, 3)).astype(np.float32)))
self.soft = P.SoftmaxCrossEntropyWithLogits()
self.depend = depend
default='',
help='maskrcnn ckpt file.')
parser.add_argument('--output_file',
type=str,
default='',
help='maskrcnn output air name.')
args_opt = parser.parse_args()
net = Mask_Rcnn_Resnet50(config=config)
param_dict = load_checkpoint(args_opt.ckpt_file)
load_param_into_net(net, param_dict)
net.set_train(False)
bs = config.test_batch_size
img = Tensor(np.zeros([bs, 3, 768, 1280], np.float16))
img_metas = Tensor(np.zeros([bs, 4], np.float16))
gt_bboxes = Tensor(np.zeros([bs, 128, 4], np.float16))
gt_labels = Tensor(np.zeros([bs, 128], np.int32))
gt_num = Tensor(np.zeros([bs, 128], np.bool))
gt_mask = Tensor(np.zeros([bs, 128], np.bool))
export(net,
img,
img_metas,
gt_bboxes,
gt_labels,
gt_num,
gt_mask,
file_name=args_opt.output_file,
file_format="AIR")
class Dihedral(nn.Cell):
"""dihedral class"""
def __init__(self, controller):
super(Dihedral, self).__init__()
self.CONSTANT_Pi = 3.1415926535897932
if controller.amber_parm is not None:
file_path = controller.amber_parm
self.read_information_from_amberfile(file_path)
self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32)
self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32)
self.atom_c = Tensor(np.asarray(self.h_atom_c, np.int32), mstype.int32)
self.atom_d = Tensor(np.asarray(self.h_atom_d, np.int32), mstype.int32)
self.pk = Tensor(np.asarray(self.pk, np.float32), mstype.float32)
self.gamc = Tensor(np.asarray(self.gamc, np.float32), mstype.float32)
self.gams = Tensor(np.asarray(self.gams, np.float32), mstype.float32)
self.pn = Tensor(np.asarray(self.pn, np.float32), mstype.float32)
self.ipn = Tensor(np.asarray(self.ipn, np.int32), mstype.int32)
def process1(self, context):
"""process1: read information from amberfile"""
for idx, val in enumerate(context):
if idx < len(context) - 1:
if "%FLAG POINTERS" in val + context[
idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]:
start_idx = idx + 2
count = 0
value = list(map(int, context[start_idx].strip().split()))
self.dihedral_with_hydrogen = value[6]
self.dihedral_numbers = value[7]
self.dihedral_numbers += self.dihedral_with_hydrogen
information = []
information.extend(value)
while count < 15:
start_idx += 1
value = list(
map(int, context[start_idx].strip().split()))
information.extend(value)
count += len(value)
self.dihedral_type_numbers = information[17]
print("dihedral type numbers ", self.dihedral_type_numbers)
break
self.phase_type = [0] * self.dihedral_type_numbers
self.pk_type = [0] * self.dihedral_type_numbers
self.pn_type = [0] * self.dihedral_type_numbers
for idx, val in enumerate(context):
if "%FLAG DIHEDRAL_FORCE_CONSTANT" in val:
count = 0
start_idx = idx
information = []
while count < self.dihedral_type_numbers:
start_idx += 1
if "%FORMAT" in context[start_idx]:
continue
else:
value = list(
map(float, context[start_idx].strip().split()))
information.extend(value)
count += len(value)
self.pk_type = information[:self.dihedral_type_numbers]
break
for idx, val in enumerate(context):
if "%FLAG DIHEDRAL_PHASE" in val:
count = 0
start_idx = idx
information = []
while count < self.dihedral_type_numbers:
start_idx += 1
if "%FORMAT" in context[start_idx]:
continue
else:
value = list(
map(float, context[start_idx].strip().split()))
information.extend(value)
count += len(value)
self.phase_type = information[:self.dihedral_type_numbers]
break
for idx, val in enumerate(context):
if "%FLAG DIHEDRAL_PERIODICITY" in val:
count = 0
start_idx = idx
information = []
while count < self.dihedral_type_numbers:
start_idx += 1
if "%FORMAT" in context[start_idx]:
continue
else:
value = list(
map(float, context[start_idx].strip().split()))
information.extend(value)
count += len(value)
self.pn_type = information[:self.dihedral_type_numbers]
break
def read_information_from_amberfile(self, file_path):
"""read information from amberfile"""
file = open(file_path, 'r')
context = file.readlines()
file.close()
self.process1(context)
self.h_atom_a = [0] * self.dihedral_numbers
self.h_atom_b = [0] * self.dihedral_numbers
self.h_atom_c = [0] * self.dihedral_numbers
self.h_atom_d = [0] * self.dihedral_numbers
self.pk = []
self.gamc = []
self.gams = []
self.pn = []
self.ipn = []
for idx, val in enumerate(context):
if "%FLAG DIHEDRALS_INC_HYDROGEN" in val:
count = 0
start_idx = idx
information = []
while count < 5 * self.dihedral_with_hydrogen:
start_idx += 1
if "%FORMAT" in context[start_idx]:
continue
else:
value = list(
map(int, context[start_idx].strip().split()))
information.extend(value)
count += len(value)
for i in range(self.dihedral_with_hydrogen):
self.h_atom_a[i] = information[i * 5 + 0] / 3
self.h_atom_b[i] = information[i * 5 + 1] / 3
self.h_atom_c[i] = information[i * 5 + 2] / 3
self.h_atom_d[i] = abs(information[i * 5 + 3] / 3)
tmpi = information[i * 5 + 4] - 1
self.pk.append(self.pk_type[tmpi])
tmpf = self.phase_type[tmpi]
if abs(tmpf - self.CONSTANT_Pi) <= 0.001:
tmpf = self.CONSTANT_Pi
tmpf2 = math.cos(tmpf)
if abs(tmpf2) < 1e-6:
tmpf2 = 0
self.gamc.append(tmpf2 * self.pk[i])
tmpf2 = math.sin(tmpf)
if abs(tmpf2) < 1e-6:
tmpf2 = 0
self.gams.append(tmpf2 * self.pk[i])
self.pn.append(abs(self.pn_type[tmpi]))
self.ipn.append(int(self.pn[i] + 0.001))
break
for idx, val in enumerate(context):
if "%FLAG DIHEDRALS_WITHOUT_HYDROGEN" in val:
count = 0
start_idx = idx
information = []
while count < 5 * (self.dihedral_numbers -
self.dihedral_with_hydrogen):
start_idx += 1
if "%FORMAT" in context[start_idx]:
continue
else:
value = list(
map(int, context[start_idx].strip().split()))
information.extend(value)
count += len(value)
for i in range(self.dihedral_with_hydrogen,
self.dihedral_numbers):
self.h_atom_a[i] = information[
(i - self.dihedral_with_hydrogen) * 5 + 0] / 3
self.h_atom_b[i] = information[
(i - self.dihedral_with_hydrogen) * 5 + 1] / 3
self.h_atom_c[i] = information[
(i - self.dihedral_with_hydrogen) * 5 + 2] / 3
self.h_atom_d[i] = abs(
information[(i - self.dihedral_with_hydrogen) * 5 + 3]
/ 3)
tmpi = information[(i - self.dihedral_with_hydrogen) * 5 +
4] - 1
self.pk.append(self.pk_type[tmpi])
tmpf = self.phase_type[tmpi]
if abs(tmpf - self.CONSTANT_Pi) <= 0.001:
tmpf = self.CONSTANT_Pi
tmpf2 = math.cos(tmpf)
if abs(tmpf2) < 1e-6:
tmpf2 = 0
self.gamc.append(tmpf2 * self.pk[i])
tmpf2 = math.sin(tmpf)
if abs(tmpf2) < 1e-6:
tmpf2 = 0
self.gams.append(tmpf2 * self.pk[i])
self.pn.append(abs(self.pn_type[tmpi]))
self.ipn.append(int(self.pn[i] + 0.001))
break
for i in range(self.dihedral_numbers):
if self.h_atom_c[i] < 0:
self.h_atom_c[i] *= -1
def Dihedral_Engergy(self, uint_crd, uint_dr_to_dr_cof):
"""compute dihedral energy"""
self.dihedral_energy = P.DihedralEnergy(self.dihedral_numbers)(
uint_crd, uint_dr_to_dr_cof, self.atom_a, self.atom_b, self.atom_c,
self.atom_d, self.ipn, self.pk, self.gamc, self.gams, self.pn)
self.sigma_of_dihedral_ene = P.ReduceSum()(self.dihedral_energy)
return self.sigma_of_dihedral_ene
def Dihedral_Force_With_Atom_Energy(self, uint_crd, scaler):
"""compute dihedral force and atom energy"""
self.dfae = P.DihedralForceWithAtomEnergy(
dihedral_numbers=self.dihedral_numbers)
self.frc, self.ene = self.dfae(uint_crd, scaler, self.atom_a,
self.atom_b, self.atom_c, self.atom_d,
self.ipn, self.pk, self.gamc, self.gams,
self.pn)
return self.frc, self.ene
def repeat_elements(x, rep, axis):
repeat_elements_net = RepeatElementsNet(rep, axis)
return repeat_elements_net(Tensor(x.astype(np.int32))).asnumpy()
def test_net(): x = np.random.randn(1, 16, 1, 1).astype(np.float16) reshape = Net() output = reshape(Tensor(x)) print(output.asnumpy())
def test_net():
logical_and = Net()
output = logical_and(Tensor(x1), Tensor(x2))
print(x1)
print(x2)
print(output.asnumpy())
def __init__(self, exclusive=False, reverse=False):
super(CumSum, self).__init__()
self.cumsum_op = P.CumSum(exclusive, reverse)
self.x0 = Tensor(x0)
self.axis0 = axis0
self.x1 = Tensor(x0)
self.axis1 = axis1
self.x2 = Tensor(x0)
self.axis2 = axis2
self.x3 = Tensor(x0)
self.axis3 = axis3
self.x4 = Tensor(x0)
self.axis4 = axis4
self.x5 = Tensor(x0)
self.axis5 = axis5
self.x6 = Tensor(x0)
self.axis6 = axis6
self.x7 = Tensor(x1)
self.axis7 = axis0
self.x8 = Tensor(x1)
self.axis8 = axis1
self.x9 = Tensor(x1)
self.axis9 = axis2
self.x10 = Tensor(x1)
self.axis10 = axis3
self.x11 = Tensor(x1)
self.axis11 = axis4
self.x12 = Tensor(x1)
self.axis12 = axis5
self.x13 = Tensor(x1)
self.axis13 = axis6
self.x14 = Tensor(x2)
self.axis14 = axis0
self.x15 = Tensor(x2)
self.axis15 = axis1
self.x16 = Tensor(x2)
self.axis16 = axis2
self.x17 = Tensor(x2)
self.axis17 = axis3
self.x18 = Tensor(x2)
self.axis18 = axis4
self.x19 = Tensor(x2)
self.axis19 = axis5
self.x20 = Tensor(x2)
self.axis20 = axis6
self.x21 = Tensor(x3)
self.axis21 = axis0
self.x22 = Tensor(x3)
self.axis22 = axis1
self.x23 = Tensor(x3)
self.axis23 = axis2
self.x24 = Tensor(x3)
self.axis24 = axis3
self.x25 = Tensor(x3)
self.axis25 = axis4
self.x26 = Tensor(x3)
self.axis26 = axis5
self.x27 = Tensor(x3)
self.axis27 = axis6
self.x28 = Tensor(x4)
self.axis28 = axis0
self.x29 = Tensor(x4)
self.axis29 = axis1
self.x30 = Tensor(x4)
self.axis30 = axis2
self.x31 = Tensor(x4)
self.axis31 = axis3
self.x32 = Tensor(x4)
self.axis32 = axis4
self.x33 = Tensor(x4)
self.axis33 = axis5
self.x34 = Tensor(x4)
self.axis34 = axis6
self.x35 = Tensor(x5)
self.axis35 = axis0
def __init__(
self,
num_atomtypes=100,
dim_atomembedding=64,
min_rbf_dis=0.05,
max_rbf_dis=1,
num_rbf=32,
n_interactions=3,
n_heads=8,
max_cycles=10,
activation=Swish(),
output_dim=1,
self_dis=None,
rbf_sigma=None,
distance_expansion=LogGaussianDistribution,
cutoff=None,
cutoff_network=SmoothCutoff,
public_filter=True,
coupled_interactions=False,
trainable_gaussians=False,
use_pondering=True,
fixed_cycles=False,
rescale_rbf=True,
use_time_embedding=True,
use_all_interactions=True,
use_mcr=False,
debug=False,
):
super().__init__(
num_atomtypes=num_atomtypes,
dim_atomembedding=dim_atomembedding,
min_rbf_dis=min_rbf_dis,
max_rbf_dis=max_rbf_dis,
num_rbf=num_rbf,
output_dim=output_dim,
rbf_sigma=rbf_sigma,
distance_expansion=distance_expansion,
cutoff=cutoff,
cutoff_network=cutoff_network,
rescale_rbf=rescale_rbf,
use_all_interactions=use_all_interactions,
)
self.network_name = 'AirNet'
self.max_distance = max_rbf_dis
self.min_distance = min_rbf_dis
if self_dis is None:
self.self_dis = self.min_distance
else:
self.self_dis = self_dis
self.self_dis_tensor = Tensor([self.self_dis], ms.float32)
self.n_heads = n_heads
if use_time_embedding:
time_embedding = self._get_time_signal(max_cycles,
dim_atomembedding)
else:
time_embedding = [0 for _ in range(max_cycles)]
if public_filter:
inter_filter = False
self.filter = Filter(num_rbf, dim_atomembedding, None)
else:
inter_filter = True
self.filter = None
self.n_interactions = n_interactions
# block for computing interaction
if coupled_interactions:
# use the same SchNetInteraction instance (hence the same weights)
self.interactions = nn.CellList([
AirNetInteraction(
dim_atom_embed=dim_atomembedding,
num_rbf=num_rbf,
n_heads=n_heads,
activation=activation,
max_cycles=max_cycles,
time_embedding=time_embedding,
use_filter=inter_filter,
use_pondering=use_pondering,
fixed_cycles=fixed_cycles,
)
] * n_interactions)
else:
# use one SchNetInteraction instance for each interaction
self.interactions = nn.CellList([
AirNetInteraction(
dim_atom_embed=dim_atomembedding,
num_rbf=num_rbf,
n_heads=n_heads,
activation=activation,
max_cycles=max_cycles,
time_embedding=time_embedding,
use_filter=inter_filter,
use_pondering=use_pondering,
fixed_cycles=fixed_cycles,
) for i in range(n_interactions)
])
# readout layer
if self.use_all_interactions and n_interactions > 1:
if use_mcr:
self.gather_interactions = MultipleChannelRepresentation(
n_interactions, dim_atomembedding, 1, activation)
else:
self.gather_interactions = TensorSum()
else:
self.gather_interactions = None
readoutdim = int(dim_atomembedding / 2)
self.readout = AtomwiseReadout(dim_atomembedding, self.output_dim, [
readoutdim,
], activation)
if debug:
self.debug_fun = self._debug_fun
self.lmax_label = []
for i in range(n_interactions):
self.lmax_label.append('l' + str(i) + '_cycles')
self.fill = P.Fill()
self.concat = P.Concat(-1)
self.pack = P.Pack(-1)
self.reducesum = P.ReduceSum()
self.reducemax = P.ReduceMax()
self.tensor_summary = P.TensorSummary()
self.scalar_summary = P.ScalarSummary()
class Net(Cell):
def __init__(self, mul_weight, strategy1=None, strategy2=None):
super(Net, self).__init__()
self.mul = P.Mul().set_strategy(strategy1)
self.square = P.Square().set_strategy(strategy2)
self.mul2 = P.Mul().set_strategy(strategy1)
self.mul_weight = Parameter(mul_weight, "w1")
def construct(self, x, b):
out = self.mul(x, self.mul_weight)
out = self.square(out)
out = self.mul2(out, b)
return out
_x = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
_w1 = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
_b = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
def compile_net(net):
optimizer = Momentum(net.trainable_params(),
learning_rate=0.1,
momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
_executor.compile(train_net, _x, _b)
context.reset_auto_parallel_context()
def test_square_data_parallel():
def __init__(
self,
model,
scale=1.0,
shift=0.0,
max_atoms_num=0,
aggregate=True,
average=False,
atom_types=None,
full_connect=False,
):
super().__init__()
self.predict = model
# dim_atomembedding=model.dim_atomembedding
self.full_connect = full_connect
self.scale = scale
self.shift = shift
self.aggregate = aggregate
self.average = average
self.reducesum = P.ReduceSum(keep_dims=False)
self.molsum = P.ReduceSum(keep_dims=True)
self.reducemean = P.ReduceMean(keep_dims=False)
if atom_types is None:
self.fixed_atoms = False
self.num_atoms = 0
else:
self.fixed_atoms = True
model._set_fixed_atoms(True)
if len(atom_types.shape) == 1:
self.num_atoms = len(atom_types)
elif len(atom_types.shape) == 2:
self.num_atoms = len(atom_types[0])
if self.num_atoms <= 0:
raise ValueError(
"The 'num_atoms' cannot be 0 " +
"'atom_types' is not 'None' in MolCalculator!")
if type(atom_types) is not Tensor:
atom_types = Tensor(atom_types, ms.int32)
self.atom_types = atom_types
self.neighbors = None
self.mask = None
self.fc_neighbors = None
if self.full_connect:
if self.fixed_atoms:
self.fc_neighbors = Types2FullConnectNeighbors(self.num_atoms)
self.neighbors = self.fc_neighbors.get_full_neighbors()
else:
if max_atoms_num <= 0:
raise ValueError(
"The 'max_atoms_num' cannot be 0 " +
"when the 'full_connect' flag is 'True' and " +
"'atom_types' is 'None' in MolCalculator!")
self.fc_neighbors = Types2FullConnectNeighbors(max_atoms_num)
if self.fixed_atoms and self.full_connect:
self.distances = AtomDistances(True)
model._set_fixed_neighbors()
else:
self.distances = AtomDistances(False)
self.ones = P.Ones()
def test_rmsprop():
learning_rate, decay, momentum, epsilon, centered = [0.5, 0.8, 0.9, 1e-3, True]
variable_np = np.array([1.0, 2.0], dtype=np.float32)
gradients_np = np.array([0.1, 0.2], dtype=np.float32)
mean_gradients_np = np.array([0.0, 0.0], dtype=np.float32)
mean_square_np = np.array([epsilon, epsilon], dtype=np.float32)
moment_np = np.array([0.0, 0.0], dtype=np.float32)
variable_ms = Tensor(variable_np)
gradients_ms = Tensor(gradients_np)
mean_gradients_ms = Tensor(mean_gradients_np)
mean_square_ms = Tensor(mean_square_np)
moment_ms = Tensor(moment_np)
if centered:
rmspropcented_numpy(variable_np, gradients_np, mean_gradients_np, mean_square_np, moment_np,
learning_rate, decay, momentum, epsilon)
else:
rmsprop_numpy(variable_np, gradients_np, mean_square_np, moment_np,
learning_rate, decay, momentum, epsilon)
net = NetRMSProp(centered)
_ = net(variable_ms, gradients_ms, mean_gradients_ms, mean_square_ms,
moment_ms, learning_rate, decay, momentum, epsilon)
error = np.ones(shape=variable_np.shape) * 10e-6
diff = variable_ms.asnumpy() - variable_np
assert np.all(diff < error)
error = np.ones(shape=gradients_np.shape) * 10e-6
diff = gradients_ms.asnumpy() - gradients_np
assert np.all(diff < error)
error = np.ones(shape=mean_gradients_np.shape) * 10e-6
diff = mean_gradients_ms.asnumpy() - mean_gradients_np
assert np.all(diff < error)
error = np.ones(shape=mean_square_np.shape) * 10e-6
diff = mean_square_ms.asnumpy() - mean_square_np
assert np.all(diff < error)
error = np.ones(shape=moment_np.shape) * 10e-6
diff = moment_ms.asnumpy() - moment_np
assert np.all(diff < error)
def test_dict3():
input_np = np.random.randn(2, 3, 4, 5).astype(np.float32)
input_me = Tensor(input_np)
net = Net3()
out_me = net(input_me)
assert out_me == ('x', 'y', 0, (0, 1))
