-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmultiple_gpus.py
213 lines (184 loc) · 8.49 KB
/
multiple_gpus.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
# Parameters and DataLoaders
input_size = 5
output_size = 2
batch_size = 30
data_size = 100
######################################################################
# Device
#
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
######################################################################
# Dummy DataSet
# -------------
#
# Make a dummy (random) dataset. You just need to implement the
# getitem
#
class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)
def __getitem__(self, index):
return self.data[index]
def __len__(self):
return self.len
rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
batch_size=batch_size, shuffle=True)
######################################################################
# Simple Model
# ------------
#
# For the demo, our model just gets an input, performs a linear operation, and
# gives an output. However, you can use ``DataParallel`` on any model (CNN, RNN,
# Capsule Net etc.)
#
# We've placed a print statement inside the model to monitor the size of input
# and output tensors.
# Please pay attention to what is printed at batch rank 0.
#
class Model(nn.Module):
# Our model
def __init__(self, input_size, output_size):
super(Model, self).__init__()
self.fc = nn.LSTM(input_size, output_size)
def forward(self, input):
output, _ = self.fc(input)
print("\tIn Model: input size", input.size(),
"output size", output.size())
return output
######################################################################
# Create Model and DataParallel
# -----------------------------
#
# This is the core part of the tutorial. First, we need to make a model instance
# and check if we have multiple GPUs. If we have multiple GPUs, we can wrap
# our model using ``nn.DataParallel``. Then we can put our model on GPUs by
# ``model.to(device)``
#
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
model.to(device)
######################################################################
# Run the Model
# -------------
#
# Now we can see the sizes of input and output tensors.
#
for data in rand_loader:
input = data.to(device)
output = model(input)
print("Outside: input size", input.size(),
"output_size", output.size())
######################################################################
# Results
# -------
#
# If you have no GPU or one GPU, when we batch 30 inputs and 30 outputs, the model gets 30 and outputs 30 as
# expected. But if you have multiple GPUs, then you can get results like this.
#
# 2 GPUs
# ~~~~~~
#
# If you have 2, you will see:
#
# .. code:: bash
#
# # on 2 GPUs
# Let's use 2 GPUs!
# In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
# In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
# In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
# In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
# In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
# Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
#
# 3 GPUs
# ~~~~~~
#
# If you have 3 GPUs, you will see:
#
# .. code:: bash
#
# Let's use 3 GPUs!
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
#
# 8 GPUs
# ~~~~~~~~~~~~~~
#
# If you have 8, you will see:
#
# .. code:: bash
#
# Let's use 8 GPUs!
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
# Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
#
######################################################################
# Summary
# -------
#
# DataParallel splits your data automatically and sends job orders to multiple
# models on several GPUs. After each model finishes their job, DataParallel
# collects and merges the results before returning it to you.
#
# For more information, please check out
# https://pytorch.org/tutorials/beginner/former\_torchies/parallelism\_tutorial.html.
#