У меня большая архитектура, которая не вписывается в память графического процессора, но есть замечательное свойство этой архитектуры, когда только подмножества архитектуры работают в любой момент времени в течение некоторого времени. Поэтому я хотел бы динамически загружать / выгружать веса слоев, которые не используются между процессором и графическим процессором. Как этого достичь?
Первое, что можно попробовать, это вызвать .cpu() или .cuda() с параметрами, которые я хочу переместить. К сожалению, , что может вызвать проблемы с тренировкой с оптимизатором, как указано в документации :
cuda(device=None)
Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.
Одним из примеров использования будет реализация ProxylessNAS , однако только окончательные обученные модели доступны на момент написания , а реализация поиска в архитектуре недоступна.