import torch
import torch.multiprocessing as mp
import torch.functional as f
import threading
import numpy as np
from timeit import timeit

def build(cuda=False):
    nn = torch.nn.Sequential(
        torch.nn.Linear(1024, 1024),
        torch.nn.Linear(1024, 1)
    )

    return nn.cuda() if cuda else nn

def train(nn, X, y, epoch=100):
    X = torch.autograd.Variable(X)
    y = torch.autograd.Variable(y)
    optim = torch.optim.SGD(nn.parameters(), lr=0.1)
    for i in range(epoch):
        yhat = nn(X)
        loss = ((yhat - y) ** 2).mean()
        loss.backward()
        optim.step()

def data(cuda=False):
    X = torch.Tensor(np.random.randn(10, 1024))
    y = torch.Tensor(np.random.randn(10, 1))
    return (X.cuda(), y.cuda()) if cuda else (X, y)

def cpu_run(i=None):
    nn = build(cuda=False)
    d = data(cuda=False)
    train(nn, *d)

def seq_cpu_run():
    for i in range(5):
        cpu_run()

def multiprocess_cpu_run():
    pool = torch.multiprocessing.Pool(processes=1)
    result = pool.map(cpu_run, [() for i in range(1)])
    pool.close()
    pool.join()
    return result

if __name__ == "__main__":
    print(timeit(seq_cpu_run, number=1)) # 1
    print(timeit(multiprocess_cpu_run, number=1))  # 2