File size: 5,214 Bytes
1d8fc13 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | #!/usr/bin/python
"""
precompute_teacher.py
Pre-computes DLoc teacher outputs for all training samples and saves them.
Run this ONCE before distillation — eliminates teacher forward pass during training.
Usage:
python precompute_teacher.py
Output:
./data/teacher_train_outputs.pt (~2.3 GB for 17K samples at [1,161,361])
./data/teacher_test_outputs.pt
"""
import os
import time
import torch
import numpy as np
import warnings
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=FutureWarning)
from utils import *
from data_loader import load_data
from modelADT import ModelADT
from joint_model import Enc_2Dec_Network
# ============================================================
# CONFIG
# ============================================================
TEACHER_RUN_DIR = './runs/2026-05-21-06-51-03'
TEACHER_EPOCH = 'best'
BATCH_SIZE = 64
trainpath = [
'./data/dataset_jacobs_July28.mat',
'./data/dataset_non_fov_train_jacobs_July28_2.mat',
'./data/dataset_fov_train_jacobs_July28_2.mat',
]
testpath = [
'./data/dataset_fov_test_jacobs_July28_2.mat',
'./data/dataset_non_fov_test_jacobs_July28_2.mat',
]
# ============================================================
# LOAD DATA
# ============================================================
device = torch.device('cuda:0')
print('Loading training data...')
train_wo, train_w, train_lbl = [], [], []
for fp in trainpath:
wo, w, lbl = load_data(fp)
if lbl.dim() == 3:
lbl = lbl.unsqueeze(1)
train_wo.append(wo)
train_w.append(w)
train_lbl.append(lbl)
train_wo = torch.cat(train_wo, dim=0)
train_w = torch.cat(train_w, dim=0)
train_lbl = torch.cat(train_lbl, dim=0)
print(f'Training samples: {train_w.shape[0]}')
print('Loading testing data...')
test_wo, test_w, test_lbl = [], [], []
for fp in testpath:
wo, w, lbl = load_data(fp)
if lbl.dim() == 3:
lbl = lbl.unsqueeze(1)
test_wo.append(wo)
test_w.append(w)
test_lbl.append(lbl)
test_wo = torch.cat(test_wo, dim=0)
test_w = torch.cat(test_w, dim=0)
test_lbl = torch.cat(test_lbl, dim=0)
print(f'Testing samples: {test_w.shape[0]}')
# ============================================================
# LOAD TEACHER
# ============================================================
print('Loading teacher model...')
from params_storage.params_fig10b_single_gpu import opt_exp, opt_encoder, opt_decoder, opt_offset_decoder
opt_encoder.checkpoints_load_dir = TEACHER_RUN_DIR
opt_decoder.checkpoints_load_dir = TEACHER_RUN_DIR
opt_offset_decoder.checkpoints_load_dir = TEACHER_RUN_DIR
enc_model = ModelADT()
enc_model.initialize(opt_encoder)
enc_model.setup(opt_encoder)
dec_model = ModelADT()
dec_model.initialize(opt_decoder)
dec_model.setup(opt_decoder)
offset_dec_model = ModelADT()
offset_dec_model.initialize(opt_offset_decoder)
offset_dec_model.setup(opt_offset_decoder)
teacher = Enc_2Dec_Network()
teacher.initialize(opt_exp, enc_model, dec_model, offset_dec_model, gpu_ids=opt_exp.gpu_ids)
enc_model.load_networks(TEACHER_EPOCH, load_dir=TEACHER_RUN_DIR)
dec_model.load_networks(TEACHER_EPOCH, load_dir=TEACHER_RUN_DIR)
offset_dec_model.load_networks(TEACHER_EPOCH, load_dir=TEACHER_RUN_DIR)
teacher.eval()
print('Teacher loaded.')
# ============================================================
# PRE-COMPUTE OUTPUTS
# ============================================================
def precompute(name, w_data, lbl_data, wo_data):
"""Run teacher on all samples, return stacked outputs on CPU."""
loader = torch.utils.data.DataLoader(
torch.utils.data.TensorDataset(wo_data, w_data, lbl_data),
batch_size=BATCH_SIZE, shuffle=False,
num_workers=4, pin_memory=True,
)
outputs = []
t0 = time.time()
with torch.no_grad():
for i, data in enumerate(loader):
teacher.set_input(data[1], data[2], data[0], shuffle_channel=False)
teacher.test()
outputs.append(teacher.decoder.output.cpu())
if (i + 1) % 50 == 0:
print(f' {name}: {(i+1)*BATCH_SIZE}/{len(w_data)} samples')
result = torch.cat(outputs, dim=0)
elapsed = time.time() - t0
print(f' {name}: {result.shape} in {elapsed:.1f}s')
return result
print('\nPre-computing teacher outputs for training set...')
train_teacher_out = precompute('train', train_w, train_lbl, train_wo)
print('Pre-computing teacher outputs for test set...')
test_teacher_out = precompute('test', test_w, test_lbl, test_wo)
# ============================================================
# SAVE
# ============================================================
train_path = './data/teacher_train_outputs.pt'
test_path = './data/teacher_test_outputs.pt'
torch.save(train_teacher_out, train_path)
torch.save(test_teacher_out, test_path)
train_size = os.path.getsize(train_path) / 1e9
test_size = os.path.getsize(test_path) / 1e9
print(f'\nSaved:')
print(f' {train_path} ({train_size:.2f} GB)')
print(f' {test_path} ({test_size:.2f} GB)')
print('Done. Now run train_distillation.py')
|