Пример оптимизатора ADAM в PyTorch
Введение
Часто полезно заморозить некоторые параметры, например, когда вы настраиваете свою модель и хотите заморозить некоторые слои в зависимости от примера, который вы обрабатываете, как показано на рисунке.
Как мы видим, в первом примере мы замораживаем первые два слоя и обновляем параметры последних двух, а во втором примере мы замораживаем второй и четвертый слои и настраиваем остальные. Будет много других случаев, когда этот метод будет полезен, и если вы читаете эту статью, у вас, вероятно, будет случай для этого.
Постановка проблемы
Чтобы немного упростить ситуацию, давайте предположим, что у нас есть модель, которая принимает два разных типа входных данных — один с 3 функциями, другой с 2 функциями, и в зависимости от того, какие входные данные передаются, мы будем передавать их через два разных начальных слоя. Таким образом, мы хотим обновить только параметры, связанные с этими конкретными входными данными во время обучения. Как видно ниже, мы хотим заморозить слой hidden_task1 при передаче input1 и заморозить слой hidden_task2 при передаче input2.
class Network(nn.Module):
def __init__(self):
super().__init__()
# Inputs to hidden layer linear transformation
self.hidden_task1 = nn.Linear(3, 3, bias=False)
self.hidden_task2 = nn.Linear(2, 3, bias=False)
self.output = nn.Linear(3, 4, bias=False)
# Define sigmoid activation and softmax output
self.sigmoid = nn.Sigmoid()
self.softmax = nn.Softmax(dim=1)
def forward(self, x, task='task1'):
if task == 'task1':
x = self.hidden_task1(x)
else:
x = self.hidden_task2(x)
x = self.sigmoid(x)
x = self.output(x)
x = self.softmax(x)
return x
def freeze_params(self, params_str):
for n, p in self.named_parameters():
if n in params_str:
p.grad = None
def freeze_params_grad(self, params_str):
for n, p in self.named_parameters():
if n in params_str:
p.requires_grad = False
def unfreeze_params_grad(self, params_str):
for n, p in self.named_parameters():
if n in params_str:
p.requires_grad = True
# define input and target
input1 = torch.randn(10, 3).to(device)
input2 = torch.randn(10, 2).to(device)
target1 = torch.randint(0, 4, (10, )).long().to(device)
target2 = torch.randint(0, 4, (10, )).long().to(device)
net = Network().to(device)
# helper
def changed_parameters(initial, final):
for n, p in initial.items():
if not torch.allclose(p, final[n]):
print("Changed : ", n)
В мире, где есть только оптимизаторы SGD
Если бы мы работали только с оптимизатором SGD, проблема была бы решена просто с помощью requires_grad = False, который не будет вычислять градиенты для указанных нами параметров, и, таким образом, мы получим желаемые результаты.
original_param = {n : p.clone() for (n, p) in net.named_parameters()}
print("Original params ")
pprint(original_param)
print(100 * "=")
# let's define 2 loss functions (we could only define one actually
# in this case as they are the same)
criterion1 = nn.CrossEntropyLoss()
criterion2 = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.9)
# set requires_grad to False for selected layers
net.freeze_params_grad(['hidden_task2.weight'])
print("Params after task 1 update ")
params_hid1 = {n : p.clone() for (n, p) in net.named_parameters()}
pprint(params_hid1)
print(100 * "=")
# output for task 1 - we want to keep frozen task2 layer parameters
output = net(input1, task='task1')
optimizer.zero_grad() # zero the gradient buffers
loss1 = criterion(output, target)
loss1.backward()
optimizer.step()
print("States optimizer 1: ")
print(optimizer.state)
# set requires_grad back to True for selected layers
net.unfreeze_params_grad(['hidden_task2.weight'])
# output for task 2 - we want to keep frozen task1 layer parameters
output1 = net(input2, task='task2')
optimizer.zero_grad() # zero the gradient buffers
loss2 = criterion1(output1, target1)
loss2.backward()
optimizer.step() # Does the update
print("States optimizer 1: ")
print(optimizer.state)
# set requires_grad back to True for selected layers
net.unfreeze_params_grad(['hidden_task1.weight'])
print("Params after task 2 update ")
params_hid2 = {n : p.clone() for (n, p) in net.named_parameters()}
pprint(params_hid2)
changed_parameters(params_hid1, params_hid2)
В выходных данных ниже мы видим, что параметр «Changed» после обновлений задачи 1 и задачи 2 правильный, и мы достигли желаемого результата.
{'hidden_task1.weight': tensor([[-0.0043, 0.3097, -0.4752],
[-0.4249, -0.2224, 0.1548],
[-0.0114, 0.4578, -0.0512]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1871, -0.2137],
[-0.1390, -0.6755],
[-0.4683, -0.2915]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.0214, 0.2282, 0.3464],
[-0.3914, -0.2514, 0.2097],
[ 0.4794, -0.1188, 0.4320],
[-0.0931, 0.0611, 0.5228]], device='cuda:0',
grad_fn=<CloneBackward0>)}
====================================================================================================
Params after hidden
{'hidden_task1.weight': tensor([[ 0.0010, 0.3107, -0.4746],
[-0.4289, -0.2261, 0.1547],
[-0.0105, 0.4596, -0.0528]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1871, -0.2137],
[-0.1390, -0.6755],
[-0.4683, -0.2915]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.0554, 0.2788, 0.3800],
[-0.4105, -0.2702, 0.1917],
[ 0.4552, -0.1496, 0.4091],
[-0.0838, 0.0601, 0.5301]], device='cuda:0',
grad_fn=<CloneBackward0>)}
====================================================================================================
Changed : hidden_task1.weight
Changed : output.weight
Params after hidden 1
{'hidden_task1.weight': tensor([[ 0.0010, 0.3107, -0.4746],
[-0.4289, -0.2261, 0.1547],
[-0.0105, 0.4596, -0.0528]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1906, -0.2102],
[-0.1412, -0.6783],
[-0.4657, -0.2929]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.0386, 0.2673, 0.3726],
[-0.3818, -0.2414, 0.2232],
[ 0.4402, -0.1698, 0.3898],
[-0.0807, 0.0631, 0.5254]], device='cuda:0',
grad_fn=<CloneBackward0>)}
Changed : hidden_task2.weight
Changed : output.weight
Сложности с адаптивными оптимизаторами
Теперь попробуем снова запустить то же самое, но с помощью оптимизатора Adam:
optimizer = optim.Adam(net.parameters(), lr=0.9)
В части «Изменено» теперь мы видим, что после второго обновления задачи hidden_task1.weight также изменился, а это не то, что нам нужно.
Original params
{'hidden_task1.weight': tensor([[-0.0043, 0.3097, -0.4752],
[-0.4249, -0.2224, 0.1548],
[-0.0114, 0.4578, -0.0512]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1871, -0.2137],
[-0.1390, -0.6755],
[-0.4683, -0.2915]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.0214, 0.2282, 0.3464],
[-0.3914, -0.2514, 0.2097],
[ 0.4794, -0.1188, 0.4320],
[-0.0931, 0.0611, 0.5228]], device='cuda:0',
grad_fn=<CloneBackward0>)}
====================================================================================================
Params after hidden
{'hidden_task1.weight': tensor([[ 0.8957, 1.2069, 0.4291],
[-1.3211, -1.1204, -0.7465],
[ 0.8887, 1.3537, -0.9508]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1871, -0.2137],
[-0.1390, -0.6755],
[-0.4683, -0.2915]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.9212, 1.1262, 1.2433],
[-1.2879, -1.1492, -0.6922],
[-0.4249, -1.0177, -0.4718],
[ 0.8078, -0.8394, 1.4181]], device='cuda:0',
grad_fn=<CloneBackward0>)}
====================================================================================================
Changed : hidden_task1.weight
Changed : output.weight
Params after hidden 1
{'hidden_task1.weight': tensor([[ 1.4907, 1.7991, 1.0283],
[-1.9122, -1.7133, -1.3428],
[ 1.4837, 1.9445, -1.5453]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[-0.7146, -1.1118],
[-1.0377, 0.2305],
[-1.3641, -1.1889]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.9372, 1.3922, 1.5032],
[-1.5886, -1.4844, -0.9789],
[-0.8855, -1.5812, -1.0326],
[ 1.6785, -0.2048, 2.3004]], device='cuda:0',
grad_fn=<CloneBackward0>)}
Changed : hidden_task1.weight
Changed : hidden_task2.weight
Changed : output.weight
Попробуем понять, что здесь происходит. Правило обновления для SGD определяется как:
Где альфа — скорость обучения, nabla L — градиент относительно параметров. Как мы видим, если градиент равен нулю, параметры не обновляются, так как правило обновления является только функцией градиентов. И когда мы установим requires_grad = False, градиенты для этих слоев будут равны нулю и не будут вычисляться.
Как насчет адаптивных оптимизаторов, таких как ADAM или других, где правило обновления является не только функцией градиентов? Давайте посмотрим на АДАМ:
Где Beta1, Beta2 — некоторые гиперпараметры, alpha — скорость обучения, mt — первый момент и vtявляется вторым моментом градиентов gt. Это правило обновления позволяет вычислить скорость адаптивного обучения для каждого параметра.
Что наиболее важно для нашей проблемы, даже если текущий градиент gt установлен равным нулю через requires_grad = False , параметры все равно обновляются оптимизатором с использованием сохраненных mt и vtзначения. В самом деле, если мы напечатаем optimizer.state, мы увидим, что оптимизатор хранит количество шагов (т. е. количество градиентныхобновлений каждого параметра), exp_avg, что является первый момент и exp_avg_sqвторой момент:
# optimizer step 1
defaultdict(<class 'dict'>, {Parameter containing:
tensor([[ 0.8957, 1.2069, 0.4291],
[-1.3211, -1.1204, -0.7465],
[ 0.8887, 1.3537, -0.9508]], device='cuda:0', requires_grad=True):
{'step': tensor(1.),
'exp_avg': tensor([[-5.9304e-04, -1.0966e-04, -5.9985e-05],
[ 4.4068e-04, 4.1636e-04, 1.7705e-05],
[-1.0544e-04, -2.0357e-04, 1.7783e-04]], device='cuda:0'),
'exp_avg_sq': tensor([[3.5170e-08, 1.2025e-09, 3.5982e-10],
[1.9420e-08, 1.7336e-08, 3.1345e-11],
[1.1118e-09, 4.1440e-09, 3.1623e-09]], device='cuda:0')},
Parameter containing:
tensor([[ 0.9212, 1.1262, 1.2433],
[-1.2879, -1.1492, -0.6922],
[-0.4249, -1.0177, -0.4718],
[ 0.8078, -0.8394, 1.4181]], device='cuda:0', requires_grad=True):
{'step': tensor(1.),
'exp_avg': tensor([[-0.0038, -0.0056, -0.0037],
[ 0.0021, 0.0021, 0.0020],
[ 0.0027, 0.0034, 0.0025],
[-0.0010, 0.0001, -0.0008]], device='cuda:0'),
'exp_avg_sq': tensor([[1.4261e-06, 3.1517e-06, 1.3953e-06],
[4.4782e-07, 4.3352e-07, 3.9994e-07],
[7.2213e-07, 1.1702e-06, 6.4754e-07],
[1.0547e-07, 1.2353e-09, 6.5470e-08]], device='cuda:0')}})
# optimizer step 2
tensor([[ 1.4907, 1.7991, 1.0283],
[-1.9122, -1.7133, -1.3428],
[ 1.4837, 1.9445, -1.5453]], device='cuda:0', requires_grad=True):
{'step': tensor(2.),
'exp_avg': tensor([[-5.3374e-04, -9.8693e-05, -5.3987e-05],
[ 3.9661e-04, 3.7472e-04, 1.5934e-05],
[-9.4899e-05, -1.8321e-04, 1.6005e-04]], device='cuda:0'),
'exp_avg_sq': tensor([[3.5135e-08, 1.2013e-09, 3.5946e-10],
[1.9400e-08, 1.7318e-08, 3.1314e-11],
[1.1107e-09, 4.1398e-09, 3.1592e-09]], device='cuda:0')},
Parameter containing:
tensor([[ 0.9372, 1.3922, 1.5032],
[-1.5886, -1.4844, -0.9789],
[-0.8855, -1.5812, -1.0326],
[ 1.6785, -0.2048, 2.3004]], device='cuda:0', requires_grad=True):
{'step': tensor(2.), 'exp_avg': tensor([[-0.0002, -0.0025, -0.0017],
[ 0.0011, 0.0011, 0.0010],
[ 0.0019, 0.0029, 0.0021],
[-0.0028, -0.0015, -0.0014]], device='cuda:0'),
'exp_avg_sq': tensor([[2.4608e-06, 3.7819e-06, 1.6833e-06],
[5.1839e-07, 4.8712e-07, 4.7173e-07],
[7.4856e-07, 1.1713e-06, 6.4888e-07],
[4.4950e-07, 2.6660e-07, 1.1588e-07]], device='cuda:0')},
Parameter containing:
tensor([[-0.7146, -1.1118],
[-1.0377, 0.2305],
[-1.3641, -1.1889]], device='cuda:0', requires_grad=True):
{'step': tensor(1.),
'exp_avg': tensor([[ 0.0009, 0.0011],
[ 0.0045, -0.0002],
[ 0.0003, 0.0012]], device='cuda:0'),
'exp_avg_sq': tensor([[8.7413e-08, 1.3188e-07],
[1.9946e-06, 4.3840e-09],
[8.1403e-09, 1.3691e-07]], device='cuda:0')}})
Мы видим, что в первом обновлении optimizer.step() мы получаем только два параметра в состояниях оптимизатора — hidden_task1и output. На втором шаге оптимизатора у нас есть все параметры, но обратите внимание, что hidden_task1 обновляется дважды, чего не должно быть.
Итак, как с ними бороться? Решение на самом деле очень простое — вместо использования набора requires_grad просто установите grad = None для параметров. Таким образом, код становится:
original_param = {n : p.clone() for (n, p) in net.named_parameters()}
print("Original params ")
pprint(original_param)
print(100 * "=")
# let's define 2 loss functions (we could only define one actually
# in this case as they are the same)
criterion1 = nn.CrossEntropyLoss()
criterion2 = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.9)
print("Params after task 1 update ")
params_hid1 = {n : p.clone() for (n, p) in net.named_parameters()}
pprint(params_hid1)
print(100 * "=")
# output for task 1 - we want to keep frozen task2 layer parameters
output = net(input1, task='task1')
optimizer.zero_grad() # zero the gradient buffers
loss1 = criterion1(output, target1)
loss1.backward()
# Freeze parameters here!
net.freeze_params(['hidden_task2.weight'])
optimizer.step()
# output for task 2 - we want to keep frozen task1 layer parameters
output = net(input2, task='task2')
optimizer.zero_grad() # zero the gradient buffers
loss2 = criterion2(output, target2)
loss2.backward()
# Freeze parameters here!
net.freeze_params_grad(['hidden_task1.weight'])
optimizer.step() # Does the update
print("Params after task 2 update ")
params_hid2 = {n : p.clone() for (n, p) in net.named_parameters()}
pprint(params_hid2)
changed_parameters(params_hid1, params_hid2)
Обратите внимание, что нам нужно установить
grad = Noneпослеloss.backward(), так как нам нужно сначала вычислить градиенты для всех параметров, но доoptimizer.step().
Если мы сейчас запустим код оптимизатора ADAM, результаты будут такими, как ожидалось.
Original params
{'hidden_task1.weight': tensor([[-0.0043, 0.3097, -0.4752],
[-0.4249, -0.2224, 0.1548],
[-0.0114, 0.4578, -0.0512]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1871, -0.2137],
[-0.1390, -0.6755],
[-0.4683, -0.2915]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.0214, 0.2282, 0.3464],
[-0.3914, -0.2514, 0.2097],
[ 0.4794, -0.1188, 0.4320],
[-0.0931, 0.0611, 0.5228]], device='cuda:0',
grad_fn=<CloneBackward0>)}
====================================================================================================
Params after task 1 update
{'hidden_task1.weight': tensor([[ 0.8957, 1.2069, 0.4291],
[-1.3211, -1.1204, -0.7465],
[ 0.8887, 1.3537, -0.9508]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[ 0.1871, -0.2137],
[-0.1390, -0.6755],
[-0.4683, -0.2915]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.9212, 1.1262, 1.2433],
[-1.2879, -1.1492, -0.6922],
[-0.4249, -1.0177, -0.4718],
[ 0.8078, -0.8394, 1.4181]], device='cuda:0',
grad_fn=<CloneBackward0>)}
====================================================================================================
Changed : hidden_task1.weight
Changed : output.weight
Params after task 2 update
{'hidden_task1.weight': tensor([[ 0.8957, 1.2069, 0.4291],
[-1.3211, -1.1204, -0.7465],
[ 0.8887, 1.3537, -0.9508]], device='cuda:0',
grad_fn=<CloneBackward0>),
'hidden_task2.weight': tensor([[-0.7146, -1.1118],
[-1.0377, 0.2305],
[-1.3641, -1.1889]], device='cuda:0', grad_fn=<CloneBackward0>),
'output.weight': tensor([[ 0.9372, 1.3922, 1.5032],
[-1.5886, -1.4844, -0.9789],
[-0.8855, -1.5812, -1.0326],
[ 1.6785, -0.2048, 2.3004]], device='cuda:0',
grad_fn=<CloneBackward0>)}
Changed : hidden_task2.weight
Changed : output.weight
Также изменился optimizer.state — на втором шаге оптимизатора hidden_task1 не обновляется и его значение step равно 1.
tensor([[ 0.8957, 1.2069, 0.4291],
[-1.3211, -1.1204, -0.7465],
[ 0.8887, 1.3537, -0.9508]], device='cuda:0', requires_grad=True):
{'step': tensor(1.),
'exp_avg': tensor([[-5.9304e-04, -1.0966e-04, -5.9985e-05],
[ 4.4068e-04, 4.1636e-04, 1.7705e-05],
[-1.0544e-04, -2.0357e-04, 1.7783e-04]], device='cuda:0'),
'exp_avg_sq': tensor([[3.5170e-08, 1.2025e-09, 3.5982e-10],
[1.9420e-08, 1.7336e-08, 3.1345e-11],
[1.1118e-09, 4.1440e-09, 3.1623e-09]], device='cuda:0')},
Parameter containing:
tensor([[ 0.9372, 1.3922, 1.5032],
[-1.5886, -1.4844, -0.9789],
[-0.8855, -1.5812, -1.0326],
[ 1.6785, -0.2048, 2.3004]], device='cuda:0', requires_grad=True):
{'step': tensor(2.),
'exp_avg': tensor([[-0.0002, -0.0025, -0.0017],
[ 0.0011, 0.0011, 0.0010],
[ 0.0019, 0.0029, 0.0021],
[-0.0028, -0.0015, -0.0014]], device='cuda:0'),
'exp_avg_sq': tensor([[2.4608e-06, 3.7819e-06, 1.6833e-06],
[5.1839e-07, 4.8712e-07, 4.7173e-07],
[7.4856e-07, 1.1713e-06, 6.4888e-07],
[4.4950e-07, 2.6660e-07, 1.1588e-07]], device='cuda:0')},
Parameter containing:
tensor([[-0.7146, -1.1118],
[-1.0377, 0.2305],
[-1.3641, -1.1889]], device='cuda:0', requires_grad=True):
{'step': tensor(1.),
'exp_avg': tensor([[ 0.0009, 0.0011],
[ 0.0045, -0.0002],
[ 0.0003, 0.0012]], device='cuda:0'),
'exp_avg_sq': tensor([[8.7413e-08, 1.3188e-07],
[1.9946e-06, 4.3840e-09],
[8.1403e-09, 1.3691e-07]], device='cuda:0')}})
Параллельное распределение данных
В качестве дополнительного примечания, если мы хотим, чтобы поддержка DistributedDataParallel в PyTorch работала с несколькими графическими процессорами, нам нужно немного изменить реализацию, описанную выше, следующим образом:
Это немного сложнее, и если вы знаете более чистый способ написать это, пожалуйста, поделитесь в комментариях!
Обратная связь
Я был бы признателен за любые отзывы о вышеизложенном — если вы знаете, могут ли возникнуть какие-либо потенциальные проблемы при выполнении этого способа и есть ли другие способы добиться того же.
Выводы
В этой статье мы описали, как сделать заморозку слоев, когда во время обучения нам нужно заморозить и разморозить некоторые слои. Если вы хотите полностью заморозить некоторые слои во время всего обучения, вы можете использовать оба решения, описанные в этой статье, так как в вашем случае не имеет значения, используете ли вы SGD или адаптивный оптимизатор. Однако, как мы видели, проблема возникает, когда вам нужно заморозить и разморозить слои во время обучения, и другое поведение, которое мы наблюдаем при использовании оптимизаторов, чье правило обновления зависит только от градиента, и тех, чье правило обновления зависит от других переменных, таких как импульс. Вы также можете найти полный код здесь.
