Не удается загрузить модель сегментации изображения частично в Pytorch - PullRequest
Новая
       66

Не удается загрузить модель сегментации изображения частично в Pytorch

0 голосов
/ 10 ноября 2019

Я пытаюсь загрузить модель частично (то есть, вместо того, чтобы загружать все слои одновременно, я просто пытаюсь загрузить первые пару слоев сети). Вот мой код: 

import torch

unet = my_unet(in_ch=5, out_ch=1).cuda()
enc = torch.nn.Sequential(*list(unet.children())[:10])
del unet # Comment this line if you wnat to load unet and see output later
print(torch.cuda.empty_cache())
x = np.random.randn(1, 5, 320, 320)
x = (x - np.min(x))/(np.max(x) - np.min(x))
out = enc(torch.FloatTensor(x).cuda())
# out = unet(torch.FloatTensor(x).cuda())
print(out.size())

Если я запускаю этот скрипт, он прерывается ошибкой CUDA out of memory. Но если я загружаю всю модель my_unet (комментируя 3-ю и 7-ю строку и раскомментируя 8-ю строку), она работает нормально и печатает out.size () (1, 1, 320, 320), как и ожидалось. Что я здесь не так делаю? Почему я могу загрузить всю модель, но не могу загрузить часть той же модели, которая требует меньше памяти? Есть ли обходной путь?

1 Ответ

0 голосов
/ 10 ноября 2019

Вы не сказали, куда брошен CUDA out of memory. Во всяком случае, я бы не стал переносить модель в графический процессор, пока у вас не появится последняя версия. Сделайте что-то вроде этого: 

unet = my_unet(in_ch=5, out_ch=1)                              # removed .cuda()
enc = torch.nn.Sequential(*list(unet.children())[:10]).cuda()  # added .cuda()

# [...]
...