Как я могу просто вычислить часть градиента, используя make_functional в PyTorch? ⇐ Python

[Anonymous]

Код: Выделить всё

func_model, func_params = make_functional(self.model)
def fm(x, func_params):
fx = func_model(func_params, x)
return fx.squeeze(0).squeeze(0)
def floss(func_params,input):
fx = fm(input, func_params)
return fx

per_sample_grads =vmap(jacrev(floss), (None, 0))(func_params, input)
cnt=0
for g in per_sample_grads:
g = g.detach()
J_d = g.reshape(len(g),-1) if cnt == 0 else torch.hstack([J_d,g.reshape(len(g),-1)])
cnt = 1

result = J_d.detach()

В этом коде per_sample_grads включает всю градацию сетевой модели путем ввода некоторых данных. Где модель

Код: Выделить всё

self.model = Network(self.input_size, self.hidden_size,  self.output_size, self.depth, act=torch.nn.Tanh() )

Я использую этот код для вычисления матрицы Якоби. [Матрица Якобиана: номер данных * номер параметра]Однако я хочу получить определенную небольшую часть градации.

Код: Выделить всё

params = torch.cat([p.view(-1) for p in self.model.parameters()], dim=0)
selected_columns = torch.random.choice(p_number, opt_num, replace=False)
target_params=params[selected_columns]# I just need to caculate the grads of the target_params

И получение всего градиента — это своего рода трата в моих кодах. Я хочу сэкономить время на этом этапе auto_grad. Что я могу сделать, чтобы добиться этого?
Как я могу вычислить градиенты только для части модели, чтобы сэкономить время?

Подробнее здесь: https://stackoverflow.com/questions/789 ... in-pytorch