Градиенты пользовательских параметров не распространяются в PyTorch ⇐ Python

[Anonymous]

Я пытаюсь реализовать слияние моделей для T5-small, где хочу узнать коэффициенты слияния во время обучения. У меня есть эталонная реализация, которая работает для других моделей, но при ее адаптации для T5-small градиенты не перетекают в коэффициенты слияния.
вот ключевая часть эталонного кода:

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

class AdaMerging(torch.nn.Module):
def __init__(self, paramslist, model, names, exam_datasets):
super(AdaMerging, self).__init__()
self.paramslist = paramslist
self.model = model
self.names = names
self.pretrain_lambdas = torch.ones(len(paramslist[0]), 1)
prior = 0.3
rlambdas = torch.ones(len(paramslist[0]), len(paramslist)-1) * prior  # (1 * tasks)
self.lambdas_raw = torch.nn.Parameter(rlambdas)

self.classifier = []
for dataset_name in exam_datasets:
classification_head = get_classification_head(args, dataset_name)
layer_name = 'classifier_{}'.format(dataset_name)
self.add_module(layer_name, classification_head.to(args.device))
self.classifier.append(layer_name)

def lambdas(self):
task_lambdas = torch.clamp(self.lambdas_raw, min=0.0, max=1.0)
lambdass = torch.cat((self.pretrain_lambdas, task_lambdas), 1)
return lambdass

def get_classification_head(self, dataset_name):
layer_name = 'classifier_{}'.format(dataset_name)
classification_head = getattr(self, layer_name)
return classification_head

def get_image_encoder(self):
alph = self.lambdas()
params = tuple(sum(tuple(pi * lambdasi for pi, lambdasi in zip(p, alph[j].cpu()))) for j, p in enumerate(zip(*self.paramslist)))
params = tuple(p.cuda(0) for p in params)
load_weights(self.model, self.names, params)
return self.model

def forward(self, inp, dataset_name):
alph = self.lambdas()
params = tuple(sum(tuple(pi * lambdasi for pi, lambdasi in zip(p, alph[j].cpu()))) for j, p in enumerate(zip(*self.paramslist)))

params = tuple(p.cuda(0) for p in params)
load_weights(self.model, self.names, params)
feature = self.model(inp)

layer_name = 'classifier_{}'.format(dataset_name)
classification_head = getattr(self, layer_name)
out = classification_head(feature)
return out

def softmax_entropy(x):
return -(x.softmax(1) * x.log_softmax(1)).sum(1)

...
...
for epoch in range(epochs):
losses = 0.
for dataset_name in exam_datasets:
dataset = get_dataset(dataset_name, pretrained_model.val_preprocess, location=args.data_location, batch_size=16)
dataloader = get_dataloader_shuffle(dataset)
for i, data in enumerate(tqdm.tqdm(dataloader)):
data = maybe_dictionarize(data)
x = data['images'].to(args.device)
y = data['labels'].to(args.device)

outputs = adamerging(x, dataset_name)
loss = softmax_entropy(outputs).mean(0)
losses += loss
optimizer.zero_grad()
losses.backward()
optimizer.step()

Я использовал Cursor для генерации кода, но он не сработал — градиенты не распространялись:

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

class ModelMerging(nn.Module):
def __init__(self, model, paramslist, names):
super(ModelMerging, self).__init__()
self.model = model
self.names = names

assert all(p1.shape == p2.shape for p1, p2 in zip(paramslist[0], paramslist[1])), \
"Model parameters must have the same shapes"

self.paramslist = paramslist

num_params = len(paramslist[0])
lambda1 = torch.ones(num_params, 1) * LAMBDA_1_INIT
lambda2 = torch.ones(num_params, 1) * LAMBDA_2_INIT
self.lambdas_raw = nn.Parameter(torch.cat([lambda1, lambda2], dim=1))

for param in self.model.parameters():
param.requires_grad = False

def forward(self, **inputs):
lambdas = self.get_lambdas()
device = lambdas.device
params = []
for j, (p1, p2) in enumerate(zip(self.paramslist[0], self.paramslist[1])):
p1 = p1.to(device)
p2 = p2.to(device)
merged_param = p1 * lambdas[j, 0] + p2 * lambdas[j, 1]
params.append(merged_param)

state_dict = {name:  param for name, param in zip(self.names, params)}

with torch.no_grad():
self.model.load_state_dict(state_dict)

outputs = self.model(**inputs)

if isinstance(outputs.loss, torch.Tensor):
outputs.loss = outputs.loss.clone()

return outputs

...
...
for epoch in range(NUM_EPOCHS):
total_loss = 0
for step, batch in enumerate(train_dataloader):
batch = {k: v.to(device) for k, v in batch.items()}

outputs = merged_model(**batch)
loss = outputs.loss

loss.backward()
optimizer.step()
optimizer.zero_grad()

total_loss += loss.item()

Мой вопрос: почему мой код не может обучать пользовательские параметры, а ссылочный код может?

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