1. 程式人生 > >pytorch 在測試階段,進行前向推斷時,視訊記憶體累加溢位問題

pytorch 在測試階段,進行前向推斷時,視訊記憶體累加溢位問題

1、問題描述:

  pytorch中,在測試階段進行前向推斷執行時,隨著for迴圈次數的增加,視訊記憶體不斷累加變大,最終導致視訊記憶體溢位。


2、解決方法:
    使用如下程式碼處理輸入資料:

  假設X為模型的輸入

  X = X.cuda()

  input_blobs = Variable(X, volatile=True)

  output = model(input_blobs)

  注意: 一定要設定 volatile=True 該引數,否則在for迴圈過程中,視訊記憶體會不斷累加。