pytorch一点点注意事项

多任务学习：分类＋回归

dataloader:

def __getitem__(self, idx):
    image, cate, valence, arousal = self.image_list[idx].rstrip().split(' ')
    img_name = os.path.join(self.root_dir,image)
    image = Image.open(img_name);
    label = int(cate)
''' 此处请注意，pytorch支持float32, '''
    valence = np.float32(valence)    
    arousal = np.float32(arousal)

    if self.transform:
        image = self.transform(image)
    # sample = {'image': image, 'label': label,'valence':valence,'arousal':arousal}
''' #此处不能用list()如：[valence, arousal],要用，np.array()。 '''
    sample = {'image': image, 'label': label,'val_aro':np.array([valence, arousal])}

 

Model:

''' #两个离散值就当做2分类的任务，但loss用MSE妥妥的 '''
self.fc_VA = nn.Linear(512 * block.expansion, 2) 

V_A = self.fc_VA(x)
x = self.fc(x)

return x, V_A

 

Q:一跑pytorch,未使用多GPU,且电脑卡主

R:忘使用下面这段话，导致在CPU上运行大程序，导致了接下来的结果。

model = torch.nn.DataParallel(model).cuda()