PyTorch 2.14의 torch.autograd.graph.node_creation_hook은 역전파 그래프의 각 노드가 만들어질 때 바로 실행됩니다. 나중에 그래프를 다시 훑지 않고, 그 시점에 메타데이터를 붙이거나 훅을 등록할 수 있습니다. 동기는 역전파 메모리를 그 메모리를 만든 앞 구간에 연결하는 일입니다. 예전에는 역전파가 끝난 뒤에야 “이 메모리가 어느 앞 연산에서 왔는지…

예전에는 역전파가 끝난 뒤에야 “이 메모리가 어느 앞 연산에서 왔는지”를 추측해야 했습니다. 오늘은 작은 모델에서 훅을 등록한 뒤, 노드가 생길 때마다 연산 이름·출력 shape를 리스트에 쌓아 보십시오. 참가자 노트에는 훅이 호출된 횟수, 남긴 필드, 훅을 끈 뒤 리스트가 비는지를 적습니다. “메모리가 줄었다”보다…

원문에서 이어서 확인하실 수 있습니다.

• 작은 선형층 한 번과 loss.backward()만 있는 경로에서 torch.autograd.graph.node_creation_hook을 등록합니다.

• 훅 안에서 노드 타입(또는 이름)과 출력 텐서 정보를 리스트에 남기고, backward가 끝난 뒤 리스트 길이가 0이 아닌지 확인합니다.

• 훅은 그래프가 만들어진 뒤가 아니라, 노드가 생기는 순간에 붙습니다. 등록 시점을 forward 전으로 두십시오.

https://daker.ai/community/autograd-node-creation-hook-attach-metadata-hooks