Machine-Learning
/
aiuNN
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

finetune_class #1

Merged
Fabel merged 96 commits from finetune_class into develop 2025-02-26 12:13:09 +00:00
Conversation 0 Commits 96 Files Changed 15 +4 -2
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit fde8bdcb6f - Show all commits

6
src/aiunn/finetune.py
View File

@ -142,11 +142,13 @@ for epoch in range(num_epochs):
with autocast(device_type=device.type): with autocast(device_type=device.type):
if use_checkpointing: if use_checkpointing:
# Wrap the forward pass with checkpointing to trade compute for memory. # Ensure the input tensor requires gradient so that checkpointing records the computation graph.
outputs = checkpoint(lambda x: model(x), low_res) low_res.requires_grad_()
outputs = checkpoint(model, low_res)
else: else:
outputs = model(low_res) outputs = model(low_res)
loss = criterion(outputs, high_res) loss = criterion(outputs, high_res)
scaler.scale(loss).backward() scaler.scale(loss).backward()
scaler.step(optimizer) scaler.step(optimizer)