Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Re-initialization of CUDA for every process
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as multiprocessing
import time
DEVICE = torch.device("cuda")
class Model1(nn.Module):
def __init__(self):
super(Model1, self).__init__()
self.fc1 = nn.Linear(1024, 512)
def forward(self, x):
x = F.tanh(self.fc1(x))
return x
class Model2(nn.Module):
def __init__(self):
super(Model2, self).__init__()
self.fc1 = nn.Linear(1024, 512)
def forward(self, x):
x = F.tanh(self.fc1(x))
return x
class Controller(nn.Module):
def __init__(self):
super(Controller, self).__init__()
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 10)
def forward(self, x):
x = F.tanh(self.fc1(x))
x = F.softmax(self.fc2(x), dim=1)
return x
class Test(multiprocessing.Process):
def __init__(self, model1, model2, controller, idx):
super(Test, self).__init__()
self.model1 = model1
self.model2 = model2
self.controller = controller
self.idx = idx
def run(self):
print("Starting: %d" % self.idx)
while True:
time.sleep(20)
def train():
model1 = Model1().to(DEVICE)
model2 = Model2().to(DEVICE)
jobs = []
for idx in range(5):
controller = Controller().to(DEVICE)
new_w = torch.randn(1024, dtype=torch.float, device=DEVICE)
controller.state_dict()['fc1.weight'].data.copy_(new_w)
new_process = Test(model1, model2, controller, idx)
jobs.append(new_process)
for p in jobs:
p.start()
for p in jobs:
p.join()
if __name__ == "__main__":
multiprocessing.set_start_method('spawn')
train()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment