Open main menu
CDOT Wiki
β
Search
Changes
← Older edit
Newer edit →
DPS921/PyTorch: Convolutional Neural Networks
16 bytes added
,
13:18, 30 November 2020
→
Parallelization Methods
class ToyModel(nn.Module):
def __init__(self):
super(ToyModel, self).__init__()
self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
self.relu = torch.nn.ReLU()
self.net2 = torch.nn.Linear(10, 5).to('cuda:1')
def forward(self, x):
x = self.relu(self.net1(x.to('cuda:0')))
return self.net2(x.to('cuda:1'))
The code is very similar to a single GPU implementation, except for the ''.to('cuda:x')'' calls, where ''cuda:0'' and ''cuda:1'' are each their own GPU.
Novell Rasam
56
edits