Class: Torch::NN::Parallel::DistributedDataParallel

Inherits:
Module
  • Object
show all
Defined in:
lib/torch/nn/parallel/distributed_data_parallel.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(mod, device_ids: nil, process_group: nil, broadcast_buffers: true) ⇒ DistributedDataParallel

Returns a new instance of DistributedDataParallel.

Raises:

  • (Torch::Error) 


7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
 
# File 'lib/torch/nn/parallel/distributed_data_parallel.rb', line 7

def initialize(mod, device_ids: nil, process_group: nil, broadcast_buffers: true)
  super()
  raise Torch::Error, "torch.distributed is not available" unless Torch::Distributed.available?

  @module = mod
  @broadcast_buffers = broadcast_buffers
  @process_group = process_group || Torch::Distributed.default_process_group
  raise Torch::Error, "Process group must be initialized before using DistributedDataParallel" unless @process_group

  @world_size = Torch::Distributed.get_world_size(@process_group)
  @rank = Torch::Distributed.get_rank(@process_group)
  @device = normalize_device(Array(device_ids).compact.first)
  move_to_device(@device) if @device

  synchronize_parameters
  @hook_handles = register_parameter_hooks
end

Instance Attribute Details

#module ⇒ Object (readonly)

Returns the value of attribute module.



5
6
7
 
# File 'lib/torch/nn/parallel/distributed_data_parallel.rb', line 5

def module
  @module
end

#process_group ⇒ Object (readonly)

Returns the value of attribute process_group.



5
6
7
 
# File 'lib/torch/nn/parallel/distributed_data_parallel.rb', line 5

def process_group
  @process_group
end

Instance Method Details

#forward(*inputs, **kwargs) ⇒ Object Also known as: call 



25
26
27
28
29
 
# File 'lib/torch/nn/parallel/distributed_data_parallel.rb', line 25

def forward(*inputs, **kwargs)
  outputs = @module.call(*move_inputs(inputs), **move_kwargs(kwargs))
  broadcast_buffers_if_needed
  outputs
end

#train(mode = true) ⇒ Object 



33
34
35
36
37
 
# File 'lib/torch/nn/parallel/distributed_data_parallel.rb', line 33

def train(mode = true)
  @module.train(mode)
  broadcast_buffers_if_needed
  self
end