一、GPU设置
咱们关心一下GPU怎么设置的吧。
在上一篇文章中的train.py中,
第126行 gpu = setup_ddp()
对于我来说,gpu=0
那么,129行,main(cfg, gpu, save_dir)的gpu就等于0了。
上到第25行,def main(cfg, gpu, save_dir):就把gpu=0传进来了
你在看46-50行, if train_cfg['DDP']: 其实咱们再配置文件中,这个DDP选项都是false。其实这个DDP就是分布式训练,需要多GPU。因此,gpu变量,由于不是DDP,所以根本没用上,白传递到main中了。
if train_cfg['DDP']:
sampler = DistributedSampler(trainset, dist.get_world_size(), dist.get_rank(), shuffle=True)
model = DDP(model, device_ids=[gpu])
else:
sampler = RandomSampler(trainset)
所以条件不满足,那么就跳到50行: sampler = RandomSampler(trainset)
这里,RandomSampler随机采样,其定义如下:
def __init__(self, data_source: Sized, replacement: bool = False,
num_samples: Optional[int] = None, generator=None) -> None:
replacement参数决定是否进行重复采样,如果replacement为True,将采用torch.randint来生成随机索引,其中索引序列是会存在重复值的,如果replacement为False,将采用torch.randperm函数来生成随机索引序列,此时序列不包含重复数值。
二、device设置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#device = torch.device(cfg['DEVICE'])
上面注释掉的, #device = torch.device(cfg['DEVICE']),才是源码。我改成了
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
你再看看44行,73、64行,就知道,作者要把数据倒到device里面了。