Spaces:
Build error
Build error
# Copyright (c) OpenMMLab. All rights reserved. | |
from itertools import chain | |
from torch.nn.parallel import DataParallel | |
from .scatter_gather import scatter_kwargs | |
class MMDataParallel(DataParallel): | |
"""The DataParallel module that supports DataContainer. | |
MMDataParallel has two main differences with PyTorch DataParallel: | |
- It supports a custom type :class:`DataContainer` which allows more | |
flexible control of input data during both GPU and CPU inference. | |
- It implement two more APIs ``train_step()`` and ``val_step()``. | |
Args: | |
module (:class:`nn.Module`): Module to be encapsulated. | |
device_ids (list[int]): Device IDS of modules to be scattered to. | |
Defaults to None when GPU is not available. | |
output_device (str | int): Device ID for output. Defaults to None. | |
dim (int): Dimension used to scatter the data. Defaults to 0. | |
""" | |
def __init__(self, *args, dim=0, **kwargs): | |
super(MMDataParallel, self).__init__(*args, dim=dim, **kwargs) | |
self.dim = dim | |
def forward(self, *inputs, **kwargs): | |
"""Override the original forward function. | |
The main difference lies in the CPU inference where the data in | |
:class:`DataContainers` will still be gathered. | |
""" | |
if not self.device_ids: | |
# We add the following line thus the module could gather and | |
# convert data containers as those in GPU inference | |
inputs, kwargs = self.scatter(inputs, kwargs, [-1]) | |
return self.module(*inputs[0], **kwargs[0]) | |
else: | |
return super().forward(*inputs, **kwargs) | |
def scatter(self, inputs, kwargs, device_ids): | |
return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim) | |
def train_step(self, *inputs, **kwargs): | |
if not self.device_ids: | |
# We add the following line thus the module could gather and | |
# convert data containers as those in GPU inference | |
inputs, kwargs = self.scatter(inputs, kwargs, [-1]) | |
return self.module.train_step(*inputs[0], **kwargs[0]) | |
assert len(self.device_ids) == 1, \ | |
('MMDataParallel only supports single GPU training, if you need to' | |
' train with multiple GPUs, please use MMDistributedDataParallel' | |
'instead.') | |
for t in chain(self.module.parameters(), self.module.buffers()): | |
if t.device != self.src_device_obj: | |
raise RuntimeError( | |
'module must have its parameters and buffers ' | |
f'on device {self.src_device_obj} (device_ids[0]) but ' | |
f'found one of them on device: {t.device}') | |
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids) | |
return self.module.train_step(*inputs[0], **kwargs[0]) | |
def val_step(self, *inputs, **kwargs): | |
if not self.device_ids: | |
# We add the following line thus the module could gather and | |
# convert data containers as those in GPU inference | |
inputs, kwargs = self.scatter(inputs, kwargs, [-1]) | |
return self.module.val_step(*inputs[0], **kwargs[0]) | |
assert len(self.device_ids) == 1, \ | |
('MMDataParallel only supports single GPU training, if you need to' | |
' train with multiple GPUs, please use MMDistributedDataParallel' | |
' instead.') | |
for t in chain(self.module.parameters(), self.module.buffers()): | |
if t.device != self.src_device_obj: | |
raise RuntimeError( | |
'module must have its parameters and buffers ' | |
f'on device {self.src_device_obj} (device_ids[0]) but ' | |
f'found one of them on device: {t.device}') | |
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids) | |
return self.module.val_step(*inputs[0], **kwargs[0]) | |