Python torch.broadcast_parameters方法代码示例

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import broadcast_parameters [as 别名]
def register_model(self, model, fp16):
        #  broadcast model parameters
        if self.node_count > 1:
            hvd.broadcast_parameters(model.state_dict(), root_rank=0)
        else:
            for param in model.parameters():
                torch.distributed.broadcast_multigpu([param], 0)

        # register hook for reduce when backpropagate
        self._parameter_names = {v: k for k, v in sorted(model.named_parameters())}
        self._handles = {}
        self._requires_update = set()
        self._grad_accs = []
        self._grad = []
        self._compression = hvd.Compression.fp16 if fp16 else hvd.Compression.none
        for p in model.parameters():
            if p.requires_grad:
                p.grad = p.data.new(p.size()).zero_()
                self._requires_update.add(p)
                p_tmp = p.expand_as(p)
                grad_acc = p_tmp.grad_fn.next_functions[0][0]
                grad_acc.register_hook(self._make_hook(p))
                self._grad_accs.append(grad_acc) 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import broadcast_parameters [as 别名]
def train_main():
    model = Net().to(device)
    # optimizer = optim.SGD(model.parameters(), lr=0.05)

    if hvd.rank() == 0:
        print(model)

    # Horovod: broadcast parameters.
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)

    # Horovod: scale learning rate by the number of GPUs.
    lr = 0.05
    optimizer = optim.SGD(model.parameters(), lr=lr * hvd.size())

    # Horovod: wrap optimizer with DistributedOptimizer.
    optimizer = hvd.DistributedOptimizer(optimizer,
                                         named_parameters=model.named_parameters())
    criterion = nn.BCELoss()

    batch_size = 25
    train_loader, train_sampler = get_train_loader(batch_size)
    validation_loader, validation_sampler = get_validation_loader(batch_size)

    log = get_tensorboard('simple_hvd')
    epochs = 50

    start_time = datetime.now()
    for epoch in range(1, epochs + 1):
        train(model, train_loader, train_sampler, criterion, optimizer, epoch, log)

        with torch.no_grad():
            if hvd.rank() == 0:
                print('\nValidation:')
            evaluate(model, validation_loader, validation_sampler, criterion, epoch, log)

    end_time = datetime.now()

    if hvd.rank() == 0:
        print('Total training time: {}.'.format(end_time - start_time))
        torch.save(model.state_dict(), model_file)
        print('Wrote model to', model_file) 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import broadcast_parameters [as 别名]
def _start_sync(self):
        hvd.broadcast_parameters(self.model.state_dict(), root_rank=0)
        hvd.broadcast_optimizer_state(self.optimizer, root_rank=0) 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import broadcast_parameters [as 别名]
def setup_horovod(model, learning_rate):
    """ Setup for Horovod usage.

    Args:
        model(MultitaskModel): The MultitaskModel object.
        learning_rate(float): Learning rate for the model.

    Returns: hvd.DistributedOptimizer: Optimizer to use for computing
    gradients and applying updates.

    """
    # Horovod: scale learning rate by the number of GPUs.
    optimizer = optim.Adam(model.parameters(), lr=learning_rate * hvd.size())

    # Horovod: broadcast parameters & optimizer state.
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    hvd.broadcast_optimizer_state(optimizer, root_rank=0)

    # Horovod: (optional) compression algorithm.
    compression = hvd.Compression.fp16

    # Horovod: wrap optimizer with DistributedOptimizer.
    optimizer = hvd.DistributedOptimizer(
        optimizer,
        named_parameters=model.named_parameters(),
        compression=compression,
    )

    return optimizer 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import broadcast_parameters [as 别名]
def horovod_train(self, model):
        # call setup after the ddp process has connected
        self.setup('fit')
        if self.is_function_implemented('setup', model):
            model.setup('fit')

        if torch.cuda.is_available() and self.on_gpu:
            # Horovod: pin GPU to local rank
            assert self.root_gpu == hvd.local_rank()
            torch.cuda.set_device(self.root_gpu)
            model.cuda(self.root_gpu)

        # avoid duplicating progress bar
        if hvd.rank() != 0 and self.progress_bar_callback is not None:
            self.progress_bar_callback.disable()

        # CHOOSE OPTIMIZER
        # allow for lr schedulers as well
        self.optimizers, self.lr_schedulers, self.optimizer_frequencies = self.init_optimizers(model)

        # Horovod: scale the learning rate by the number of workers to account for
        # increased total batch size
        for optimizer in self.optimizers:
            for param_group in optimizer.param_groups:
                param_group['lr'] *= hvd.size()

        if self.use_amp:
            # An example
            model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)
            self.optimizers = optimizers
            self.reinit_scheduler_properties(self.optimizers, self.lr_schedulers)

        # Horovod: broadcast parameters & optimizer state to ensure consistent initialization
        hvd.broadcast_parameters(model.state_dict(), root_rank=0)
        for optimizer in self.optimizers:
            hvd.broadcast_optimizer_state(optimizer, root_rank=0)

        def filter_named_parameters(model, optimizer):
            opt_params = set([p for group in optimizer.param_groups for p in group.get('params', [])])
            return [(name, p) for name, p in model.named_parameters() if p in opt_params]

        # Horovod: wrap optimizers to perform gradient aggregation via allreduce
        self.optimizers = [
            hvd.DistributedOptimizer(optimizer, named_parameters=filter_named_parameters(model, optimizer))
            for optimizer in self.optimizers
        ]

        # Update logger rank info from Horovod to avoid race conditions from  different ranks
        # creating directories / writing files in the same locations.
        self.global_rank = hvd.rank()
        rank_zero_only.rank = self.global_rank

        with ExitStack() as stack:
            for optimizer in self.optimizers:
                # Synchronization will be performed explicitly following backward()
                stack.enter_context(optimizer.skip_synchronize())

            self.run_pretrain_routine(model)

        # Make sure all workers have finished training before returning to the user
        hvd.join()