PyTorch 训练加速

PyTorch Dataloader 加速

 

参考源码：

https://github.com/NVIDIA/apex/blob/f5cd5ae937f168c763985f627bbf850648ea5f3f/examples/imagenet/main_amp.py#L256

class data_prefetcher():    def __init__(self, loader):        self.loader = iter(loader)        self.stream = torch.cuda.Stream()        self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)        self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)        # With Amp, it isn't necessary to manually convert data to half.        # if args.fp16:        #     self.mean = self.mean.half()        #     self.std = self.std.half()        self.preload()
    def preload(self):        try:            self.next_input, self.next_target = next(sel