Время вывода зависит от разных графических процессоров, использующих Torch

Question

Я получаю ошибку при запуске приведенного ниже кода логического вывода. В функции признать (), это займет 0,4 с до окончательного sh прогнозирования. Требуется еще 3 секунды, чтобы вернуть результат preds_str функции вызывающей стороны. Я обнаружил, что если я установлю gpu_id=0 в файле конфигурации, он мгновенно возвращается. Как я могу исправить эту ошибку? Заранее спасибо.

def recognize(imgs, model, demo_loader):

          t = time()
          model.eval()
          with torch.no_grad():
              for image_tensors, image_path_list in demo_loader:
                    batch_size = image_tensors.size(0)
                    image = image_tensors.to(config.device)
                    # For max length prediction
                    length_for_pred = torch.IntTensor([config.batch_max_length] * batch_size).to(config.device)
                    text_for_pred = torch.LongTensor(batch_size, config.batch_max_length + 1).fill_(0).to(config.device)

                    preds = model(image, text_for_pred, is_train=False)
                    _, preds_index = preds.max(2)
                    preds_str = converter.decode(preds_index, length_for_pred)

           print('time elapsed before return:'time()-t) #0.4s
           return preds_str
def main():                        
     model = Model()
     self.model.cuda(config.device)
     model = torch.nn.DataParallel(model, device_ids=[config.device], output_device=[config.device]).to(config.device)
     model.load_state_dict(torch.load(config.saved_model, map_location=config.device))
     AlignCollate_demo = AlignCollate(imgH=config.imgH, imgW=config.imgW, keep_ratio_with_pad=config.PAD)
     imgs_dataset = ImageDataset(imgs)
     demo_loader = torch.utils.data.DataLoader(imgs_dataset, batch_size=config.batch_size,shuffle=False,num_workers=int(config.workers),collate_fn=AlignCollate_demo, pin_memory=True)     
     start_time = time()
     # imgs = [img1, img2, ....]
     preds_str = recognize(imgs, model, demo_loader)
     print('time elapsed after return', time()-start_time) #3.4s

Файл конфигурации:

    class ConfigWordRecognizer:
        gpu_id = 1 #troublesome line here
        device = torch.device('cuda:{}'.format(gpu_id) if torch.cuda.is_available() else 'cpu')
        imgH = 32
        imgW = 100
        batch_size = 80
        workers = 8
        batch_max_length = 25

Answer 1

Я нашел решение из этой записи . Я установил CUDA_VISIBLE_DEVICES=1, gpu_id=0. Затем я удаляю

model = torch.nn.DataParallel(model, device_ids=[config.device], output_device=[config.device]).to(config.device)

и изменяю

model.load_state_dict(torch.load(config.saved_model, map_location=config.device))

на

model.load_state_dict(self.copyStateDict(torch.load(self.config.saved_model, map_location=self.config.device)))

Копировать функцию stateDict:

def copyStateDict(self, state_dict):
        if list(state_dict.keys())[0].startswith("module"):
            start_idx = 1
        else:
            start_idx = 0
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            name = ".".join(k.split(".")[start_idx:])
            new_state_dict[name] = v
        return new_state_dict

Модель работает хорошо на gpu1. Но я все еще не понимаю, почему, если я установлю 'gpu_id = 0', он хорошо работает на gpu0 без copyStateDict