训练一个epoch约为一个半小时以上，而两块P40需要约40分钟，和了解到的算力不符

自己安装的镜像为昇腾镜像仓库中ascend-mindspore，具体版本如下图：

之前安装服务器的时候工作人员已经安装过另外两个版本的镜像：

运行任务时服务器的状态为：（不知道为什么Memory-Usage一直不高）

网络为wavenet

下面是改写的训练代码

import mindspore

from md_model1 import Wave

import mindspore as md

import mindspore.dataset as ds

from mindspore import nn

from mindspore.train import Model, CheckpointConfig, ModelCheckpoint, LossMonitor

# 模型的基本参数

class ModelConfig:

input_channels: int = 2

residual_layers: int = 10

residual_channels: int = 64

dilation_cycle_length: int = 10

cfg = ModelConfig()

# 设置一些超参

lr = 0.0003

batch_size = 1

maxpatience = 100

maxepochs = 1000

log_batch = 2

# 数据准备

mindrecord_files = ["md.mindrecord"] # contains 1 or multiple MindRecord files

dataset = ds.MindDataset(dataset_files=mindrecord_files) # 已经自带shuffle

train_dataset, test_dataset = dataset.split([0.8, 0.2])

train_loader = train_dataset.batch(batch_size)

test_loader = test_dataset.batch(batch_size)

# 实例化网络

model = Wave(cfg)

# 损失函数和优化器

loss_fn = nn.MSELoss(reduction="mean")

optimizer = nn.Adam(params=model.trainable_params(), learning_rate=lr)

# 定义前向传播过程

def forward_fn(train_mix, train_soi):

logits = model(train_mix)

loss = loss_fn(logits, train_soi)

return loss

# 实例化梯度话函数

grad_fn = md.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)

def train_step(train_mix, train_soi):

loss, grads = grad_fn(train_mix, train_soi)

optimizer(grads)

return loss

def train(model, train_dataset):

size = train_dataset.get_dataset_size()

model.set_train()

for batch, (train_mix, train_soi) in enumerate(train_dataset.create_tuple_iterator()):

loss = train_step(train_mix, train_soi)

if batch % log_batch == 0:

loss, current = loss.asnumpy(), batch

print(f"loss: {loss:>7f} [{current:>3d}/{size:>3d}]")

def test(model, dataset, loss_fn):

num_batches = dataset.get_dataset_size()

model.set_train(False)

total, test_loss, correct = 0, 0, 0

for test_mix, test_soi in dataset.create_tuple_iterator():

pred = model(test_mix)

test_loss += loss_fn(pred, test_soi).asnumpy()

test_loss /= num_batches

print(f"Test: \n Avg MSE loss: {test_loss:>8f} \n")

return test_loss

epochs = maxepochs

min_loss = float('inf')

patience = 0

for t in range(epochs):

print(f"Epoch {t + 1}\n-------------------------------")

train(model, train_loader)

current_loss = test(model, test_loader, loss_fn)

if current_loss < min_loss:

min_loss = current_loss

else:

patience += 1

if patience > maxpatience:

mindspore.save_checkpoint(model, f"Wave_epoch_{t+1}_loss_{current_loss:.4f}.ckpt.ckpt")

print("model has been saved")

break