如何使用多处理同时运行两个预训练的 PyTorch 模型?

时间:2020-12-27 15:46:20

标签: python pytorch

我有两个预训练模型,我想使用相同的数据集进行预测。 是否可以使用多处理并行化进程。

我目前正在处理的代码:

from torch.multiprocessing import Pool, Process, set_start_method

set_start_method('spawn', force=True)

output1parallel = None
output2parallel = None

def getOutput1():
  print("IM HERE 1")
  global loader, resnet18, output1parallel
  for i,batch in enumerate(loader):
    currentBatch = batch.cuda()
    resnet18 = resnet18.cuda()
    output1parallel = resnet18(currentBatch).cpu()

    del currentBatch
def getOutput2():
  print("IM HERE 2")
  global loader, densenet,output, output2parallel
  for i,batch in enumerate(loader):
    currentBatch = batch.cuda()
    densenet = densenet.cuda()
    output2parallel = densenet(currentBatch).cpu()
    del currentBatch

start = time.time()
p1 = Process(target = getOutput1 )
p2 = Process(target = getOutput2 )

p1.start() 
p2.start() 

p1.join()
p2.join()

print("Time for parallel implementation: {}".format(time.time() - start))

此处不调用函数。

0 个答案:

没有答案
相关问题