Le framework pytorch teste rapidement s'il y a un problème avec la structure de votre modèle et obtient la forme de sortie

ps: J'étais très déprimé ces deux derniers jours, après avoir changé de modèle, l'ordinateur s'est bloqué dès que le code a fonctionné. Je vais essayer de voir si le modèle que j'ai écrit est faux. Ajoutez simplement le code suivant après le modèle:

#测试模型的代码
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 () est mon modèle entier, mais les résultats sont les suivants:

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

Mais il est resté bloqué pendant environ 20 minutes et il s'est un peu écrasé.

Plus tard, indépendamment de l'ajout du GPU spécifié:

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

Ou obtenez le GPU automatiquement:

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

Ne fonctionnera pas. Plus tard, je me suis demandé s'il était nécessaire de spécifier les variables à placer dans le GPU et le modèle à placer dans le GPU, puis je l'ai vérifié sur Internet. Modifiez ensuite le code de test comme suit:

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

Terminé en un instant et obtenu des résultats.

Ce problème est résolu, mais mon modèle se fige toujours lorsque le code s'exécute.