Why does the pytorch .backward() method occupies two more CPU threads when I want to restrict it to use only one CPU thread?

Viewed 165

I am trying to deploy my service using CPU and pytorch. I am asked to restrict the single-process service to use only one CPU thread. I have assigned torch.set_num_interop_threads(1) and torch.set_num_threads(1). The question is that when I check the actual CPU thread number using pstree <pid> I got three. Later I found that the problem happened in a.backward(). If I delete this line, I got 1 threads; if I keep this line, I got 3 CPU threads. Why did it happen and what should I do to keep the CPU thread to be one?

I wrote a simple python3 script to reproduce the problem:

import torch
import time
while(1):
    w = torch.Tensor([2])
    w = torch.autograd.Variable(w, requires_grad=True)
    x = torch.rand([1])
    x = torch.autograd.Variable(x, requires_grad=True)
    y = x**w
    y.backward()
    time.sleep(3)

Run it and use pstree <pid> to check the CPU threads, I got python───2*[{python}], but if I delete y.backward() I got only python. I believe the extra two CPU threads in this script is in the same problem with my service code.

0 Answers
Related