def adam_compile(loss_scale=1.0):
inputs = Tensor(np.ones([15, 1]).astype(np.float32))
label = Tensor(np.zeros([15, 1]).astype(np.float32))
net = Net(1, 1)
loss = MSELoss()
optimizer = Adam(net.trainable_params(),
learning_rate=1e-3,
beta1=0.9,
beta2=0.999,
eps=1e-8,
use_locking=False,
use_nesterov=False,
weight_decay=0.0,
loss_scale=loss_scale)
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepWithLossScaleCell(net_with_loss,
optimizer,
scale_sense=Tensor(
np.full((1), 1.0),
dtype=mstype.float32))
train_network.set_train()
output = train_network(inputs, label)
print("the result is ", output)
def test_compile_fp16_overflow():
inputs = Tensor(np.ones([16, 16]).astype(np.float32))
label = Tensor(np.zeros([16, 16]).astype(np.float32))
scaling_sens = Tensor(np.full((1), np.finfo(np.float32).max), dtype=mstype.float32)
net = NetFP16(16, 16)
loss = MSELoss()
optimizer = Lamb(net.trainable_params(), decay_steps=10, warmup_steps=5)
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepWithLossScaleCell(net_with_loss, optimizer)
train_network.set_train()
output = train_network(inputs, label, scaling_sens)
print("the result is ", output)
def test_momentum_compile():
inputs = Tensor(np.ones([15, 1]).astype(np.float32))
label = Tensor(np.zeros([15, 1]).astype(np.float32))
scaling_sens = Tensor(np.full((1), 1.0), dtype=mstype.float32)
net = Net(1, 1)
loss = MSELoss()
optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepWithLossScaleCell(net_with_loss, optimizer)
train_network.set_train()
output = train_network(inputs, label, scaling_sens)
print("the result is ", output)
def test_compile_fp16_lr_overflow():
inputs = Tensor(np.ones([16, 16]).astype(np.float32))
label = Tensor(np.zeros([16, 16]).astype(np.float32))
scaling_sens = Tensor(np.full((1), np.finfo(np.float32).max), dtype=mstype.float32)
lr = Tensor(np.ones([1], np.float32) * 0.1)
net = NetFP16(16, 16)
loss = MSELoss()
optimizer = Momentum(net.trainable_params(), learning_rate=lr, momentum=0.9)
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepWithLossScaleCell(net_with_loss, optimizer)
train_network.set_train()
output = train_network(inputs, label, scaling_sens)
print("the result is ", output)
def test_compile_fp16_lr_overflow_dynamic_graph():
inputs = Tensor(np.ones([16, 16]).astype(np.float32))
label = Tensor(np.zeros([16, 16]).astype(np.float32))
lr = Tensor(np.ones([1], np.float32) * 0.1)
net = NetFP16(16, 16)
loss = MSELoss()
optimizer = Momentum(net.trainable_params(), learning_rate=lr, momentum=0.9)
net_with_loss = WithLossCell(net, loss)
scale_manager = DynamicLossScaleManager()
update_cell = scale_manager.get_update_cell()
train_network = TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_update_cell=update_cell)
train_network.set_train()
output = train_network(inputs, label)
print("the result is ", output)
def test_compile_grad_error():
inputs = Tensor(np.ones([16, 16]).astype(np.float32))
label = Tensor(np.zeros([16, 16]).astype(np.float32))
lr = Tensor(np.ones([1], np.float32) * 0.1)
net = NetFP16(16, 16)
loss = MSELoss()
optimizer = Momentum(net.trainable_params(), learning_rate=lr, momentum=0.9)
net_with_loss = WithLossCell(net, loss)
scale_manager = DynamicLossScaleManager()
update_cell = scale_manager.get_update_cell()
train_network = TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_update_cell=update_cell)
train_network.set_train()
with pytest.raises(TypeError) as e:
train_network(inputs, label)
print(e)