Precisão mista automática usando PyTorch
Introdução
Modelos maiores de aprendizagem profunda precisam de mais poder computacional e recursos de memória. O treinamento mais rápido de redes neurais profundas foi alcançado através do desenvolvimento de novas técnicas. Em vez de FP32 (formato de números de ponto flutuante de precisão total), você pode usar FP16 (formato de números de ponto flutuante de meia precisão), e os pesquisadores descobriram que usá-los em conjunto é uma opção melhor.
A precisão mista permite o treinamento de meia precisão, preservando grande parte da precisão da rede de precisão simples. O termo “técnica de precisão mista” refere-se ao fato de que este método faz uso de representações de precisão simples e de meia precisão.
Nesta visão geral do treinamento automático de precisão mista (Amp) com PyTorch, demonstramos como a técnica funciona, percorrendo passo a passo o processo de uso do Amp, e discutimos aplicações mais avançadas de técnicas Amp com andaimes de código para os usuários integrarem posteriormente com seu próprio código.
Pré-requisitos
Conhecimento básico de PyTorch: Familiaridade com PyTorch, incluindo seus conceitos básicos como tensores, módulos e loop de treinamento.
Compreensão dos fundamentos do aprendizado profundo: conceitos como redes neurais, retropropagação e otimização.
Conhecimento do treinamento de precisão mista: Consciência dos benefícios e desvantagens do treinamento de precisão mista, incluindo uso reduzido de memória e computação mais rápida.
Acesso a hardware compatívele: uma GPU que suporta precisão mista, como GPUs NVIDIA com Tensor Cores (por exemplo, arquiteturas Volta, Turing, Ampere).
Configuração Python e CUDA: um ambiente Python funcional com PyTorch instalado e CUDA configurado para aceleração de GPU.
Visão geral da precisão mista
Como a maioria das estruturas de aprendizado profundo, o PyTorch normalmente treina em dados de ponto flutuante de 32 bits (FP32). O FP32, por outro lado, nem sempre é necessário para o sucesso. É possível usar ponto flutuante de 16 bits para algumas operações, onde o FP32 consome mais tempo e memória.
Consequentemente, os engenheiros da NVIDIA desenvolveram uma técnica que permite que o treinamento de precisão mista seja realizado no FP32 para um pequeno número de operações, enquanto a maior parte da rede é executada no FP16.
- Converta o modelo para utilizar o tipo de dados float16 sempre que possível.
- Manter os pesos mestres float32 para acumular atualizações de peso a cada iteração.
- O uso de escala de perda para preservar pequenos valores de gradiente.
Precisão mista em PyTorch
Para treinamento de precisão mista, o PyTorch oferece diversos recursos já integrados.
Os parâmetros de um módulo são convertidos para FP16 quando você chama o método .half(), e os dados de um tensor são convertidos para FP16 quando você chama .half(). A aritmética rápida FP16 será usada para executar quaisquer operações nesses módulos ou tensores. As bibliotecas matemáticas NVIDIA (cuBLAS e cuDNN) são bem suportadas pelo PyTorch. Os dados do pipeline FP16 são processados usando Tensor Cores para conduzir GEMMs e convoluções. Para empregar Tensor Cores em cuBLAS, as dimensões de um GEMM ([M, K] x [K, N] -> [M, N]) devem ser múltiplos de 8.
Apresentando o Apex
Os utilitários de precisão mista do Apex têm como objetivo aumentar a velocidade do treinamento, mantendo a precisão e a estabilidade do treinamento de precisão simples. O Apex pode realizar operações em FP16 ou FP32, lidar automaticamente com a conversão de parâmetros mestre e dimensionar automaticamente as perdas.
O Apex foi criado para facilitar aos pesquisadores a inclusão de treinamento de precisão mista em seus modelos. Amp, abreviação de Automatic Mixed-Precision, é um dos recursos do Apex, uma extensão leve do PyTorch. Mais algumas linhas em suas redes são suficientes para que os usuários se beneficiem do treinamento de precisão mista com Amp. O Apex foi lançado no CVPR 2018 e é importante notar que a comunidade PyTorch tem demonstrado forte apoio ao Apex desde o seu lançamento.
Ao fazer apenas pequenas alterações no modelo em execução, o Amp faz com que você não precise se preocupar com tipos mistos ao criar ou executar seu script. As suposições de Amp podem não se encaixar tão bem em modelos que utilizam PyTorch de maneiras menos usuais, mas existem ganchos para ajustar essas suposições conforme necessário.
Amp oferece todas as vantagens do treinamento de precisão mista sem a necessidade de gerenciamento explícito de escalonamento de perdas ou conversões de tipo. O site GitHub da Apex contém instruções para o processo de instalação e sua documentação oficial da API pode ser encontrada aqui.
Como funcionam os amplificadores
Amp utiliza um paradigma de lista branca/lista negra no nível lógico. As operações de tensor no PyTorch incluem funções de rede neural, como torch.nn.funcional.conv2d, funções matemáticas simples, como torch.log, e métodos de tensor, como torch.Tensor. adicionar__ . Existem três categorias principais de funções neste universo:
- Whitelist: Funções que poderiam se beneficiar do aumento de velocidade matemática do FP16. As aplicações típicas incluem multiplicação e convolução de matrizes.
- Lista Negra: As entradas devem estar em FP32 para funções onde 16 bits de precisão podem não ser suficientes.
- Todo o resto (quaisquer funções que sobrarem): Funções que podem ser executadas no FP16, mas o gasto de uma conversão de FP32 -> FP16 para executá-las no FP16 não vale a pena, pois a aceleração não é significativa .
A tarefa de Amp é simples, pelo menos em teoria. Amp determina se uma função PyTorch está na lista de permissões, na lista negra ou nenhuma delas antes de chamá-la. Todos os argumentos devem ser convertidos para FP16 se estiver na lista de permissões ou FP32 se estiver na lista negra. Caso contrário, apenas certifique-se de que todos os argumentos sejam do mesmo tipo. Esta política não é tão simples de aplicar na realidade como parece.
Usando Amp em conjunto com um modelo PyTorch
Para incluir Amp em um script PyTorch atual, siga estas etapas:
- Use a biblioteca Apex para importar Amp.
- Inicialize o Amp para que ele possa fazer as alterações necessárias no modelo, no otimizador e nas funções internas do PyTorch.
- Observe onde a retropropagação (.backward()) ocorre para que o Amp possa simultaneamente dimensionar a perda e limpar o estado por iteração.
Etapa 1
Há apenas uma linha de código para a primeira etapa:
from apex import amp
Etapa 2
O modelo de rede neural e o otimizador usado para treinamento já devem estar especificados para concluir esta etapa, que tem apenas uma linha.
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
Configurações adicionais permitem ajustar o tensor do Amp e os ajustes do tipo de operação. A função amp.initialize() aceita vários parâmetros, entre os quais especificaremos apenas três:
- modelos (torch.nn.Module ou lista de torch.nn.Modules) – Modelos para modificar/converter.
- otimizadores (opcional, torch.optim.Optimizer ou lista de torch.optim.Optimizers) – Otimizadores para modificar/cast. OBRIGATÓRIO para treinamento, opcional para inferência.
- opt_level (str, opcional, padrão=“O1 ”) – Nível de otimização de precisão puro ou misto. Os valores aceitos são “O0”, “O1”, “O2” e “O3”, explicados detalhadamente acima. Existem quatro níveis de otimização:
O0 para treinamento FP32: Este é um ambiente autônomo. Não há necessidade de se preocupar com isso, pois seu modelo de entrada já deve ser FP32 e O0 pode ajudar a estabelecer uma linha de base para precisão.
O1 para precisão mista (recomendado para uso típico): Modifique todos os métodos Tensor e Torch para usar um esquema de conversão de entrada de lista branca-lista negra. No FP16, são realizadas operações de lista de permissões, como operações amigáveis ao Tensor Core, como GEMMs e convoluções. Softmax, por exemplo, é uma operação de lista negra que requer precisão FP32. Salvo indicação explícita em contrário, O1 também emprega escalonamento de perda dinâmica.
O2 para precisão mista “Quase FP16”: O2 lança os pesos do modelo para FP16, corrige o método forward do modelo para converter dados de entrada para FP16, mantém batchnorms em FP32, mantém pesos mestres FP32, atualiza os param_groups do otimizador para que o otimizador.step() atue diretamente nos pesos FP32 e implemente a escala dinâmica de perdas (a menos que seja substituído). Ao contrário do O1, o O2 não corrige funções do Torch ou métodos do Tensor.
O3 para treinamento FP16: O3 pode não ser tão estável quanto O1 e O2 em relação à verdadeira precisão mista. Consequentemente, pode ser vantajoso definir uma velocidade de base para o seu modelo, contra a qual a eficiência de O1 e O2 possa ser avaliada.
A substituição extra da propriedade keep_batchnorm_fp32=True no O3 pode ajudá-lo a determinar a “velocidade da luz” se o seu modelo empregar normalização de lote, habilitando cudnn batchnorm.
O0 e O3 não são verdadeiras precisão mista, mas ajudam a definir linhas de base de precisão e velocidade, respectivamente. Uma implementação de precisão mista é definida como O1 e O2.
Você pode experimentar os dois e ver qual melhora mais o desempenho e a precisão para o seu modelo específico.
Etapa 3
Certifique-se de identificar onde ocorre a passagem para trás no seu código.
Algumas linhas de código semelhantes a estas aparecerão:
loss = criterion(…)
loss.backward()
optimizer.step()
Etapa 4
Usando o gerenciador de contexto Amp, você pode ativar o escalonamento de perdas simplesmente agrupando a passagem para trás:
loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
optimizer.step()
Isso é tudo. Agora você pode executar novamente seu script com o treinamento de precisão mista ativado.
Capturando chamadas de função
PyTorch não possui o objeto de modelo estático ou gráfico para travar e inserir as conversões mencionadas acima, uma vez que é muito flexível e dinâmico. Ao “monkey patching” as funções necessárias, o Amp pode interceptar e lançar parâmetros dinamicamente.
Como exemplo, você pode usar o código abaixo para garantir que os argumentos do método torch.nn.funcional.linear sejam sempre convertidos para fp16:
orig_linear = torch.nn.functional.linear
def wrapped_linear(*args):
casted_args = []
for arg in args:
if torch.is_tensor(arg) and torch.is_floating_point(arg):
casted_args.append(torch.cast(arg, torch.float16))
else:
casted_args.append(arg)
return orig_linear(*casted_args)
torch.nn.functional.linear = wrapped_linear
Embora Amp possa adicionar refinamentos para tornar o código mais resiliente, chamar Amp.init() efetivamente faz com que patches de macaco sejam inseridos em todas as funções PyTorch relevantes para que os argumentos sejam lançados corretamente em tempo de execução.
Minimizando conversões
Cada peso é convertido apenas em FP32 -> FP16 uma vez a cada iteração, pois o Amp mantém um cache interno de todas as conversões de parâmetros e os reutiliza conforme necessário. A cada iteração, o gerenciador de contexto para a passagem para trás indica ao Amp quando limpar o cache.
Autocasting e escala de gradiente usando PyTorch
“Treinamento automatizado de precisão mista” refere-se à combinação de torch.cuda.amp.autocast e torch.cuda.amp.GradScaler. Usando torch.cuda.amp.autocast, você pode configurar a transmissão automática apenas para determinadas áreas. O Autocasting seleciona automaticamente a precisão das operações da GPU para otimizar a eficiência e manter a precisão.
As instâncias torch.cuda.amp.GradScaler facilitam a execução das etapas de escala de gradiente. A escala de gradiente reduz o underflow de gradiente, o que ajuda as redes com gradientes float16 a obter melhor convergência.
Aqui está um código para demonstrar como usar autocast() para obter precisão mista automatizada no PyTorch:
# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
Creates a GradScaler once at the beginning of training.
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# Runs the forward pass with autocasting.
with autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
# Backward ops run in the same dtype autocast chose for corresponding forward ops.
scaler.scale(loss).backward()
# scaler.step() first unscales the gradients of the optimizer's assigned params.
scaler.step(optimizer)
# Updates the scale for next iteration.
scaler.update()
Se a passagem para frente para uma operação específica tiver entradas float16, então a passagem para trás para esta operação produzirá gradientes float16 e float16 poderá não ser capaz de expressar gradiente com pequenas magnitudes.
A atualização dos parâmetros relacionados será perdida se esses valores forem zerados (“underflow”).
O escalonamento gradiente é uma técnica que usa um fator de escala para multiplicar as perdas da rede e, em seguida, executa uma passagem reversa na perda escalonada para evitar estouro negativo. Também é necessário dimensionar os gradientes de fluxo reverso através da rede por esse mesmo fator. Conseqüentemente, os valores do gradiente têm uma magnitude maior, o que os impede de chegar a zero.
Antes de atualizar os parâmetros, o gradiente de cada parâmetro (atributo .grad) deve ser descalado para que o fator de escala não interfira na taxa de aprendizagem. Tanto o autocast quanto o GradScaler podem ser usados independentemente, pois são modulares.
Trabalhando com gradientes não dimensionados
Recorte gradiente
Podemos dimensionar todos os gradientes usando o método Scaler.scale(Loss).backward(). As propriedades .grad dos parâmetros entre backward() e scaler.step(optimizer) devem ser redimensionadas antes de serem alteradas ou inspecionadas. Se você deseja limitar a norma global (consulte torch.nn.utils.clip_grad_norm_()) ou a magnitude máxima (consulte torch.nn.utils.clip_grad_value_()) do seu gradiente definido como menor ou igual a um determinado valor ( algum limite imposto pelo usuário), você pode usar uma técnica chamada “recorte de gradiente. ”
Recortar sem desdimensionar resultaria na escala das magnitudes norma/máxima dos gradientes, invalidando o limite solicitado (que deveria ser o limite para gradientes não escalonados). Os gradientes contidos nos parâmetros fornecidos pelo otimizador não são dimensionados por scaler.unscale (optimizer).
Você pode remover a escala dos gradientes de outros parâmetros que foram fornecidos anteriormente a outro otimizador (como o otimizador1) usando scaler.unscale (optimizer1). Podemos ilustrar esse conceito adicionando duas linhas de códigos:
# Unscales the gradients of optimizer's assigned params in-place
scaler.unscale_(optimizer)Since the gradients of optimizer's assigned params are unscaled, clips as usual:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
Trabalhando com gradientes em escala
Acumulação de gradiente
A acumulação de gradiente é baseada em um conceito absurdamente básico. Em vez de atualizar os parâmetros do modelo, ele espera e acumula os gradientes em lotes sucessivos para calcular a perda e o gradiente.
Após um determinado número de lotes, os parâmetros são atualizados dependendo do gradiente cumulativo. Aqui está um trecho de código sobre como usar a acumulação de gradiente usando pytorch:
scaler = GradScaler()
for epoch in epochs:
for i, (input, target) in enumerate(data):
with autocast():
output = model(input)
loss = loss_fn(output, target)
# normalize the loss
loss = loss / iters_to_accumulate
# Accumulates scaled gradients.
scaler.scale(loss).backward()
# weights update
if (i + 1) % iters_to_accumulate == 0:
# may unscale_ here if desired
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
- O acúmulo de gradiente adiciona gradientes em um tamanho de lote adequado de batch_ per_iter * iters_to_accumulate.
A balança deverá estar calibrada para o lote efetivo; isso envolve verificar as notas inf/NaN, pular a etapa se algum inf/NaN for detectado e atualizar a escala para a granularidade do lote efetivo.
Também é vital manter os graduados em um fator de escala escalonado e consistente quando os graduados de um determinado lote efetivo são somados.
Se os graduados não forem escalonados (ou o fator de escala mudar) antes da acumulação ser concluída, a próxima passagem para trás adicionará graduados escalonados a graduados não escalonados (ou graduados escalados por um fator diferente), após o qual será impossível recuperar a etapa acumulada de graduados não escalonados.
- Você pode cancelar a escala dos graduados usando unscale pouco antes da etapa, após todos os graduados escalados para a próxima etapa terem sido acumulados.
Para garantir um lote completo e eficaz, basta chamar update no final de cada iteração onde você chamou anteriormente step - A função enumerate(data) nos permite acompanhar o índice do lote enquanto iteramos pelos dados.
- Divida a perda operacional por iters_to_accumulate(loss/iters_to_accumulate). Isso reduz a contribuição de cada minilote que processamos, normalizando a perda. Se você calcular a média da perda dentro de cada lote, a divisão já estará correta e nenhuma normalização adicional será necessária. Esta etapa pode não ser necessária dependendo de como você calcula a perda.
- Quando usamos
scaler.scale(loss).backward(), o PyTorch acumula os gradientes escalonados e os armazena até chamarmosoptimizer.zero grad().
Penalidade de gradiente
Ao implementar uma penalidade de gradiente, torch.autograd.grad() é usado para construir gradientes, que são combinados para formar o valor da penalidade e depois adicionados à perda. A penalidade L2 sem escalonamento ou lançamento automático é mostrada no exemplo abaixo.
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
# Creates gradients
grad_prams = torch.autograd.grad(outputs=loss,
inputs=model.parameters(),
create_graph=True)
# Computes the penalty term and adds it to the loss
grad_norm = 0
for grad in grad_prams:
grad_norm += grad.pow(2).sum()
grad_norm = grad_norm.sqrt()
loss = loss + grad_norm
loss.backward()
# You can clip gradients here
optimizer.step()
Os tensores fornecidos para torch.autograd.grad() devem ser escalonados para implementar uma penalidade de gradiente. É necessário desescalar os gradientes antes de combiná-los para obter o valor da penalidade. Como o cálculo do termo de penalidade faz parte do avanço, ele deve ocorrer dentro de um contexto de autocast.
Para a mesma penalidade L2, é assim que fica:
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast():
output = model(input)
loss = loss_fn(output, target)
# Perform loss scaling for autograd.grad's backward pass, resulting #scaled_grad_prams
scaled_grad_prams = torch.autograd.grad(outputs=scaler.scale(loss),
inputs=model.parameters(),
create_graph=True)
# Creates grad_prams before computing the penalty(grad_prams must be #unscaled).
# Because no optimizer owns scaled_grad_prams, conventional division #is used instead of scaler.unscale_:
inv_scaled = 1./scaler.get_scale()
grad_prams = [p * inv_scaled for p in scaled_grad_prams]
# The penalty term is computed and added to the loss.
with autocast():
grad_norm = 0
for grad in grad_prams:
grad_norm += grad.pow(2).sum()
grad_norm = grad_norm.sqrt()
loss = loss + grad_norm
# Applies scaling to the backward call.
# Accumulates properly scaled leaf gradients.
scaler.scale(loss).backward()
# You can unscale_ here
# step() and update() proceed as usual.
scaler.step(optimizer)
scaler.update()
Trabalhando com vários modelos, perdas e otimizadores
Scaler.scale deve ser chamado em cada perda na sua rede se você tiver muitas delas.
Se você tiver muitos otimizadores em sua rede, poderá executar scaler.unscale em qualquer um deles e deverá chamar scaler.step em cada um deles. Porém, scaler.update deve ser usado apenas uma vez, após a revisão de todos os otimizadores usados nesta iteração:
scaler = torch.cuda.amp.GradScaler()
for epoch in epochs:
for input, target in data:
optimizer1.zero_grad()
optimizer2.zero_grad()
with autocast():
output1 = model1(input)
output2 = model2(input)
loss1 = loss_fn(2 * output1 + 3 * output2, target)
loss2 = loss_fn(3 * output1 - 5 * output2, target)
#Although retain graph is unrelated to amp, it is present in this #example since both backward() calls share certain regions of graph.
scaler.scale(loss1).backward(retain_graph=True)
scaler.scale(loss2).backward()
# If you wish to view or adjust the gradients of the params they #possess, you may specify which optimizers get explicit unscaling. .
scaler.unscale_(optimizer1)
scaler.step(optimizer1)
scaler.step(optimizer2)
scaler.update()
Cada otimizador examina seus gradientes para infs/NaNs e faz um julgamento individual se deve ou não pular a etapa. Alguns otimizadores podem pular a etapa, enquanto outros não. O salto de etapas ocorre apenas uma vez a cada centenas de iterações; portanto, não deve afetar a convergência. Para modelos com vários otimizadores, você pode relatar o problema se observar uma convergência ruim após adicionar a escala de gradiente.
Trabalhando com várias GPUs
Um dos problemas mais significativos dos modelos de Deep Learning é que eles estão crescendo demais para serem treinados em uma única GPU. Pode levar muito tempo para treinar um modelo em uma única GPU, e o treinamento multi-GPU é necessário para preparar os modelos o mais rápido possível. Um conhecido pesquisador conseguiu reduzir o período de treinamento do ImageNet de duas semanas para 18 minutos ou treinar o mais extenso e avançado Transformer-XL em duas semanas, em vez de quatro anos.
DataParallel e DistributedDataParallel
Sem comprometer a qualidade, o PyTorch oferece a melhor combinação de facilidade de uso e controle. nn.DataParallel e nn.parallel.DistributedDataParallel são dois recursos do PyTorch para distribuir treinamento em várias GPUs. Você pode usar esses wrappers e alterações fáceis de usar para treinar a rede em várias GPUs.
DataParallel em um único processo
Em uma única máquina, o DataParallel ajuda a distribuir o treinamento por muitas GPUs.
Vamos dar uma olhada mais de perto em como o DataParallel realmente funciona na prática.
Ao utilizar DataParallel para treinar uma rede neural, ocorrem os seguintes estágios:
Fonte :
- O minilote é dividido em GPU:0.
- Divida e distribua lotes mínimos para todas as GPUs disponíveis.
- Copie o modelo para as GPUs.
- A passagem direta ocorre em todas as GPUs.
- Perda computacional em relação às saídas de rede na GPU:0, bem como perdas de retorno para as diversas GPUs. Os gradientes devem ser calculados em cada GPU.
- Soma dos gradientes na GPU:0 e aplique o otimizador para atualizar o modelo.
Vale a pena notar que as preocupações discutidas aqui se aplicam apenas ao autocast. O comportamento do GradScaler permanece inalterado. Não importa se torch.nn.DataParallel cria threads para cada dispositivo para conduzir a passagem direta. O estado de autocast é comunicado em cada um, e o seguinte funcionará:
model = Model_m()
p_model = nn.DataParallel(model)
Sets autocast in the main thread
with autocast():
# There will be autocasting in p_model.
output = p_model(input)
# loss_fn also autocast
loss = loss_fn(output)
DistributedDataParallel, uma GPU por processo
A documentação de torch.nn.parallel.DistributedDataParallel recomenda o uso de uma GPU por processo para melhor desempenho. Nessa situação, DistributedDataParallel não inicia threads internamente; portanto, o uso de autocast e GradScaler não é afetado.
DistributedDataParallel, múltiplas GPUs por processo
Aqui, torch.nn.parallel.DistributedDataParallel pode gerar um thread lateral para executar a passagem direta em cada dispositivo, como torch.nn.DataParallel. A correção é a mesma: aplique autocast como parte do método forward do seu modelo para garantir que ele esteja habilitado em threads laterais.
Conclusão
Neste artigo, temos:
- Apresentado o Apex.
- Vi como funcionam os amplificadores.
- Vi como realizar escalonamento de gradiente, recorte de gradiente, acumulação de gradiente e penalidade de gradinet.
- Vimos como podemos trabalhar com múltiplos modelos, perdas e otimizadores.
- Vimos como podemos executar DataParallel em um único processo ao trabalhar com várias GPUs.
Referências
https://developer.nvidia.com/blog/apex-pytorch-easy-mixed-precision-training/
https://nvidia.github.io/apex/amp.html
https://discuss.pytorch.org/t/accumulating-gradients/30020 https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2