A estrutura pytorch testa rapidamente se há um problema com a estrutura do seu modelo e obtém a forma de saída

ps: Fiquei muito deprimido nos últimos dois dias.Depois de mudar o modelo, o computador travou assim que o código rodou. Vou tentar ver se o modelo que escrevi está errado. Basta adicionar o seguinte código após o 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 () é todo o meu modelo, mas os resultados são os seguintes:

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

Mas ele travou por cerca de 20 minutos e travou um pouco.

Posteriormente, independentemente de adicionar a GPU especificada:

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

Ou obtenha a GPU automaticamente:

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')

Não funciona. Mais tarde, me perguntei se era necessário especificar as variáveis ​​a serem colocadas na GPU e o modelo a ser colocado na GPU, então verifiquei na Internet. Em seguida, altere o código de teste para o seguinte:

#测试模型的代码
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)

Concluído em um instante e obteve resultados.

Este problema foi resolvido, mas meu modelo ainda congela quando o código é executado.