|
|
| """
|
| 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
|
|
|
|
|
|
|
|
|
| TEACHER_RUN_DIR = './weights'
|
| 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',
|
| ]
|
|
|
|
|
|
|
|
|
| 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]}')
|
|
|
|
|
|
|
|
|
| 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.')
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
|
|
| 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')
|
|
|
