PyTorch и NVIdia Flare забирают все вычислительные ресурсы на эксперименты по машинному обучению ⇐ Python

[Anonymous]

Я использую PyTorch для объединенных экспериментов. Поскольку мои эксперименты включают 50 наборов данных с моделями, мне приходится запускать несколько экспериментов с моделями ML параллельно.
Код для обучения модели ML доступен здесь:

Код: Выделить всё

def train(dataloader, model, loss_fn, optimizer, device):
num_batches = len(dataloader) # Total number of observation divided by batch size
model.train()
model.to(device)

total_loss = 0
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device) # x is covariates and y is the pseudo values in the batch
# Compute prediction error
pred = model(X)
loss = loss_fn(pred,y)

# Backpropagation
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
total_loss += float(loss.item())
total_loss /= num_batches
return total_loss

Как вы можете видеть, тензор x, y и модель PyTorch переносятся на устройство cuda:0. Однако этот процесс по-прежнему выполняется на полную мощность процессора.
Я уже пробовал настраивать эти конфигурации:

Код: Выделить всё

    torch.set_num_threads(2)

Кроме того, на стороне NVidia Flare я ограничил использование 5 потоков. Тем не менее, процессор на сервере используется на 100%.

Код: Выделить всё

 controller = FederatedAvg(
patience=args.patience,
modeldir=args.modeldir,
reg=args.reg,
lr=args.lr,
optimizer=args.optimizer,
epochs=args.epochs,
partition=args.partition,
batch_size=args.batch_size,
model=args.model,
device=args.device,
dataset=args.dataset,
num_clients=args.n_parties,
num_rounds=args.comm_round,
seed=args.init_seed,
logdir=args.logdir,
run_id=args.run_id,
arguments=args,
)

Код: Выделить всё

   job.simulator_run(ws, n_clients=args.n_parties, threads=5, log_config= os.path.abspath(config_path))

В результате, когда я одновременно запускаю несколько методов ML, из-за использования 100% ЦП одновременно обрабатывается только один процесс, а другие запущенные методы ждут завершения этого эксперимента. Однако мне приходится запускать несколько методов одновременно.
Как лучше всего проводить эксперименты быстрее? Как я могу улучшиться?

Подробнее здесь: https://stackoverflow.com/questions/798 ... earning-ex