El marco de pytorch prueba rápidamente si hay un problema con la estructura de su modelo y obtiene la forma de salida

PD: Estuve muy deprimido durante los últimos dos días. Después de cambiar el modelo, la computadora se atascó tan pronto como se ejecutó el código. Intentaré ver si el modelo que escribí es incorrecto. Simplemente agregue el siguiente código después del modelo:

#测试模型的代码
def FPN():
    return Net()
def test():
    from torch.autograd import Variable
    net = FPN()
    fms = net(Variable(torch.randn(4, 1, 128, 128, 128)))
    for fm in fms:
        print(fm.size())
test()

Net () es mi modelo completo, pero los resultados son los siguientes:

(4L, 8L, 8L, 8L, 3L, 5L)
(4L, 16L, 16L, 16L, 3L, 5L)
(4L, 32L, 32L, 32L, 3L, 5L)

Pero se atascó durante unos 20 minutos y se estrelló un poco.

Más tarde, independientemente de agregar la GPU especificada:

import 
os.environ ['CUDA_VISIBLE_DEVICES'] = '0'

O obtén la GPU automáticamente:

def getFreeId():
    import pynvml

    pynvml.nvmlInit()
    def getFreeRatio(id):
        handle = pynvml.nvmlDeviceGetHandleByIndex(id)
        use = pynvml.nvmlDeviceGetUtilizationRates(handle)
        ratio = 0.5*(float(use.gpu+float(use.memory)))
        return ratio

    deviceCount = pynvml.nvmlDeviceGetCount()
    available = []
    for i in range(deviceCount):
        if getFreeRatio(i)<70:
            available.append(i)
    gpus = ''
    for g in available:
        gpus = gpus+str(g)+','
    gpus = gpus[:-1]
    return gpus
def setgpu(gpuinput):
    freeids = getFreeId()
    if gpuinput=='all':
        gpus = freeids
    else:
        gpus = gpuinput
        if any([g not in freeids for g in gpus.split(',')]):
            raise ValueError('gpu'+g+'is being used')
    print('using gpu '+gpus)
    os.environ['CUDA_VISIBLE_DEVICES']=gpus
    return len(gpus.split(','))
n_gpu = setgpu('all')

No trabajará. Luego me pregunté si era necesario especificar las variables a colocar en la GPU y el modelo a colocar en la GPU, luego lo comprobé en Internet. Luego, cambie el código de prueba a lo siguiente:

#测试模型的代码
def FPN():
    return Net()
def test():
    from torch.autograd import Variable
    net = FPN()
    net.cuda()
    fms = net(Variable(torch.randn(4, 1, 128, 128, 128).cuda(async=True)))

    for fm in fms:
        print(fm.size())
import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
test()

(4L, 8L, 8L, 8L, 3L, 5L)
(4L, 16L, 16L, 16L, 3L, 5L)
(4L, 32L, 32L, 32L, 3L, 5L)

Completado en un instante y obtenido resultados.

Este problema está resuelto, pero mi modelo aún se congela cuando se ejecuta el código.