Setting CUDA_VISIBLE_DEVICES after importing torch.utils.tensorboard doesn't work

1

Hi,

I tried to use torch.utils.tensorboard (for replacing tensorboardX) and found that:

import os

import torch
import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import DataLoader
os.environ['CUDA_VISIBLE_DEVICES'] = '4'
from torch.utils.tensorboard import SummaryWriter

(main code)

works normally, while

import os

import torch
import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
os.environ['CUDA_VISIBLE_DEVICES'] = '4'

(main code)

will still use gpu 0.

Am I using it wrong?

Thanks!

2

I would recommend to set the environment variables in the terminal before running the script, as setting it via os.environ might have side effects.
E.g. if you are loading a library, which calls to CUDA and grabs all available GPUs, setting 'CUDA_VISIBLE_DEVICES' won’t have any effect anymore, as it apparently the case in your code.
If you really want to set it in your code, make sure os.environ is called first in your script.

3

Thank you!

It makes sense. However, I am wondering why replacing torch.utils.tensorboard to tensorboardX will make the code work?

I usually set os.environ after args are parsed since gpu is specified in args, which is typically after pakages are imported.

4

I’m not sure why this happens, but I would guess that some imports in the __init__ method checks for GPUs and thus grabs them (which is apparently not the case in tensorboardX).
However, I’m also not sure it’s worth investigating. :wink: