CoCalc -- data_parallel

CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

GitHub Repository: pytorch/tutorials
Path: blob/main/beginner_source/blitz/data_parallel_tutorial.py
Views: ⁴⁹⁴
1
"""
2
Optional: Data Parallelism
3
==========================
4
**Authors**: `Sung Kim <https://github.com/hunkim>`_ and `Jenny Kang <https://github.com/jennykang>`_
5

6
In this tutorial, we will learn how to use multiple GPUs using ``DataParallel``.
7

8
It's very easy to use GPUs with PyTorch. You can put the model on a GPU:
9

10
.. code:: python
11

12
    device = torch.device("cuda:0")
13
    model.to(device)
14

15
Then, you can copy all your tensors to the GPU:
16

17
.. code:: python
18

19
    mytensor = my_tensor.to(device)
20

21
Please note that just calling ``my_tensor.to(device)`` returns a new copy of
22
``my_tensor`` on GPU instead of rewriting ``my_tensor``. You need to assign it to
23
a new tensor and use that tensor on the GPU.
24

25
It's natural to execute your forward, backward propagations on multiple GPUs.
26
However, Pytorch will only use one GPU by default. You can easily run your
27
operations on multiple GPUs by making your model run parallelly using
28
``DataParallel``:
29

30
.. code:: python
31

32
    model = nn.DataParallel(model)
33

34
That's the core behind this tutorial. We will explore it in more detail below.
35
"""
36

37

38
######################################################################
39
# Imports and parameters
40
# ----------------------
41
#
42
# Import PyTorch modules and define parameters.
43
#
44

45
import torch
46
import torch.nn as nn
47
from torch.utils.data import Dataset, DataLoader
48

49
# Parameters and DataLoaders
50
input_size = 5
51
output_size = 2
52

53
batch_size = 30
54
data_size = 100
55

56

57
######################################################################
58
# Device
59
#
60
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
61

62
######################################################################
63
# Dummy DataSet
64
# -------------
65
#
66
# Make a dummy (random) dataset. You just need to implement the
67
# getitem
68
#
69

70
class RandomDataset(Dataset):
71

72
    def __init__(self, size, length):
73
        self.len = length
74
        self.data = torch.randn(length, size)
75

76
    def __getitem__(self, index):
77
        return self.data[index]
78

79
    def __len__(self):
80
        return self.len
81

82
rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
83
                         batch_size=batch_size, shuffle=True)
84

85

86
######################################################################
87
# Simple Model
88
# ------------
89
#
90
# For the demo, our model just gets an input, performs a linear operation, and
91
# gives an output. However, you can use ``DataParallel`` on any model (CNN, RNN,
92
# Capsule Net etc.)
93
#
94
# We've placed a print statement inside the model to monitor the size of input
95
# and output tensors.
96
# Please pay attention to what is printed at batch rank 0.
97
#
98

99
class Model(nn.Module):
100
    # Our model
101

102
    def __init__(self, input_size, output_size):
103
        super(Model, self).__init__()
104
        self.fc = nn.Linear(input_size, output_size)
105

106
    def forward(self, input):
107
        output = self.fc(input)
108
        print("\tIn Model: input size", input.size(),
109
              "output size", output.size())
110

111
        return output
112

113

114
######################################################################
115
# Create Model and DataParallel
116
# -----------------------------
117
#
118
# This is the core part of the tutorial. First, we need to make a model instance
119
# and check if we have multiple GPUs. If we have multiple GPUs, we can wrap
120
# our model using ``nn.DataParallel``. Then we can put our model on GPUs by
121
# ``model.to(device)``
122
#
123

124
model = Model(input_size, output_size)
125
if torch.cuda.device_count() > 1:
126
  print("Let's use", torch.cuda.device_count(), "GPUs!")
127
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
128
  model = nn.DataParallel(model)
129

130
model.to(device)
131

132

133
######################################################################
134
# Run the Model
135
# -------------
136
#
137
# Now we can see the sizes of input and output tensors.
138
#
139

140
for data in rand_loader:
141
    input = data.to(device)
142
    output = model(input)
143
    print("Outside: input size", input.size(),
144
          "output_size", output.size())
145

146

147
######################################################################
148
# Results
149
# -------
150
#
151
# If you have no GPU or one GPU, when we batch 30 inputs and 30 outputs, the model gets 30 and outputs 30 as
152
# expected. But if you have multiple GPUs, then you can get results like this.
153
#
154
# 2 GPUs
155
# ~~~~~~
156
#
157
# If you have 2, you will see:
158
#
159
# .. code:: bash
160
#
161
#     # on 2 GPUs
162
#     Let's use 2 GPUs!
163
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
164
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
165
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
166
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
167
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
168
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
169
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
170
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
171
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
172
#         In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
173
#         In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
174
#     Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
175
#
176
# 3 GPUs
177
# ~~~~~~
178
#
179
# If you have 3 GPUs, you will see:
180
#
181
# .. code:: bash
182
#
183
#     Let's use 3 GPUs!
184
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
185
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
186
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
187
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
188
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
189
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
190
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
191
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
192
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
193
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
194
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
195
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
196
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
197
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
198
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
199
#     Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
200
#
201
# 8 GPUs
202
# ~~~~~~~~~~~~~~
203
#
204
# If you have 8, you will see:
205
#
206
# .. code:: bash
207
#
208
#     Let's use 8 GPUs!
209
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
210
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
211
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
212
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
213
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
214
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
215
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
216
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
217
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
218
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
219
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
220
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
221
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
222
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
223
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
224
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
225
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
226
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
227
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
228
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
229
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
230
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
231
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
232
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
233
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
234
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
235
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
236
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
237
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
238
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
239
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
240
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
241
#     Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
242
#
243

244

245
######################################################################
246
# Summary
247
# -------
248
#
249
# DataParallel splits your data automatically and sends job orders to multiple
250
# models on several GPUs. After each model finishes their job, DataParallel
251
# collects and merges the results before returning it to you.
252
#
253
# For more information, please check out
254
# https://pytorch.org/tutorials/beginner/former\_torchies/parallelism\_tutorial.html.
255
#
256

257
CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

Product

Resources

Company
