dloc-code / scripts /train_distillation.py

Upload scripts/train_distillation.py with huggingface_hub

778cc2b verified 5 months ago

8.5 kB

#!/usr/bin/python
"""
train_distillation.py
Knowledge distillation: DLoc teacher -> lightweight student.

Loss: L = ALPHA * MSE(student, ground_truth) + BETA * MSE(student, teacher_output)
Default: ALPHA=0.7, BETA=0.3 (from professor's spec)

Requires pre-computed teacher outputs (run precompute_teacher.py first).
Set MODE='standalone' to train without distillation (ground truth only).

Usage:
  python precompute_teacher.py          # run once
  python train_distillation.py          # distillation (default: mobilenet)

  To change student or mode, edit CONFIG below.
"""
import sys
import os
import time
import torch
import torch.nn as nn
import numpy as np
import warnings
with warnings.catch_warnings():
    warnings.filterwarnings("ignore", category=FutureWarning)

from utils import localization_error, write_log
from data_loader import load_data
from student_mobilenet import MobileNetStudent
from student_tinycnn import TinyCNNStudent

# ============================================================
# CONFIG — edit these before running
# ============================================================
STUDENT_TYPE = 'tinycnn'
MODE = 'distill'

ALPHA = 0.7                         # ground truth loss weight
BETA = 0.3                          # distillation loss weight (ignored in standalone)
EPOCHS = 50
BATCH_SIZE = 32
LR = 1e-4
NUM_WORKERS = 4

TEACHER_TRAIN_PT = './data/teacher_train_outputs.pt'
TEACHER_TEST_PT = './data/teacher_test_outputs.pt'

# ============================================================
# DATA PATHS (Fig 10b — Jacobs Hall)
# ============================================================
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',
]

# ============================================================
# SETUP
# ============================================================
torch.manual_seed(0)
np.random.seed(0)

device = torch.device('cuda:0')
save_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
run_label = f'{MODE}_{STUDENT_TYPE}'
save_dir = os.path.join('./runs', f'{run_label}_{save_name}')
os.makedirs(save_dir, exist_ok=True)

print(f'Mode: {MODE}')
print(f'Student: {STUDENT_TYPE}')
if MODE == 'distill':
    print(f'Loss: {ALPHA}*L_GT + {BETA}*L_KD')
else:
    print(f'Loss: L_GT (standalone, no distillation)')
print(f'Save dir: {save_dir}')

# ============================================================
# LOAD DATA INTO RAM
# ============================================================
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_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_w = torch.cat(test_w, dim=0)
test_lbl = torch.cat(test_lbl, dim=0)
print(f'Testing samples: {test_w.shape[0]}')

# Load pre-computed teacher outputs if distilling
if MODE == 'distill':
    print('Loading pre-computed teacher outputs...')
    train_teacher = torch.load(TEACHER_TRAIN_PT, weights_only=True)
    test_teacher = torch.load(TEACHER_TEST_PT, weights_only=True)
    assert train_teacher.shape[0] == train_w.shape[0], \
        f'Teacher/data mismatch: {train_teacher.shape[0]} vs {train_w.shape[0]}'
    print(f'Teacher outputs loaded: train={train_teacher.shape}, test={test_teacher.shape}')
    train_data = torch.utils.data.TensorDataset(train_w, train_lbl, train_teacher)
    test_data = torch.utils.data.TensorDataset(test_w, test_lbl, test_teacher)
else:
    train_data = torch.utils.data.TensorDataset(train_w, train_lbl)
    test_data = torch.utils.data.TensorDataset(test_w, test_lbl)

train_loader = torch.utils.data.DataLoader(
    train_data, batch_size=BATCH_SIZE, shuffle=True,
    num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True,
)
test_loader = torch.utils.data.DataLoader(
    test_data, batch_size=BATCH_SIZE, shuffle=False,
    num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True,
)

# ============================================================
# CREATE STUDENT
# ============================================================
if STUDENT_TYPE == 'mobilenet':
    student = MobileNetStudent(input_nc=4).to(device)
elif STUDENT_TYPE == 'tinycnn':
    student = TinyCNNStudent(input_nc=4).to(device)
else:
    raise ValueError(f'Unknown student type: {STUDENT_TYPE}')

n_params = sum(p.numel() for p in student.parameters())
print(f'Student params: {n_params:,} ({n_params/1e6:.2f}M)')

optimizer = torch.optim.Adam(student.parameters(), lr=LR)
criterion = nn.MSELoss()

# ============================================================
# TRAINING LOOP
# ============================================================
print(f'Starting {MODE} training...')
best_median = float('inf')

for epoch in range(1, EPOCHS + 1):
    epoch_start = time.time()

    # --- TRAIN ---
    student.train()
    train_loss_sum = 0
    train_errors = []

    for data in train_loader:
        inputs = data[0].to(device)
        labels = data[1].to(device)

        student_output = student(inputs)
        L_GT = criterion(student_output, labels)

        if MODE == 'distill':
            teacher_out = data[2].to(device)
            L_KD = criterion(student_output, teacher_out)
            loss = ALPHA * L_GT + BETA * L_KD
        else:
            loss = L_GT

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss_sum += loss.item()
        train_errors.extend(localization_error(
            student_output.detach().cpu().numpy(),
            labels.cpu().numpy(),
            scale=0.1,
        ))

    train_loss_avg = train_loss_sum / len(train_loader)
    train_median = np.median(train_errors)

    # --- TEST ---
    student.eval()
    test_loss_sum = 0
    test_errors = []

    with torch.no_grad():
        for data in test_loader:
            inputs = data[0].to(device)
            labels = data[1].to(device)

            student_output = student(inputs)
            loss = criterion(student_output, labels)

            test_loss_sum += loss.item()
            test_errors.extend(localization_error(
                student_output.cpu().numpy(),
                labels.cpu().numpy(),
                scale=0.1,
            ))

    test_loss_avg = test_loss_sum / len(test_loader)
    test_median = np.median(test_errors)
    test_p90 = np.percentile(test_errors, 90)
    test_p99 = np.percentile(test_errors, 99)
    elapsed = time.time() - epoch_start

    # --- LOG ---
    write_log([str(train_loss_avg)], run_label, log_dir=save_dir, log_type='train_loss')
    write_log([str(train_median)], run_label, log_dir=save_dir, log_type='train_median_error')
    write_log([str(test_loss_avg)], run_label, log_dir=save_dir, log_type='test_loss')
    write_log([str(test_median)], run_label, log_dir=save_dir, log_type='test_median_error')
    write_log([str(test_p90)], run_label, log_dir=save_dir, log_type='test_90_error')
    write_log([str(test_p99)], run_label, log_dir=save_dir, log_type='test_99_error')

    print(f'Epoch {epoch}/{EPOCHS} | {elapsed:.0f}s | '
          f'train_loss={train_loss_avg:.6f} train_med={train_median:.3f}m | '
          f'test_med={test_median:.3f}m P90={test_p90:.3f}m P99={test_p99:.3f}m')

    # --- SAVE ---
    if test_median < best_median:
        best_median = test_median
        torch.save(student.state_dict(), os.path.join(save_dir, f'best_{STUDENT_TYPE}.pth'))
        print(f'  -> New best: {best_median:.4f}m')

    torch.save(student.state_dict(), os.path.join(save_dir, f'latest_{STUDENT_TYPE}.pth'))

print(f'\nDone. Best median error: {best_median:.4f}m')
print(f'Results saved to: {save_dir}')

	#!/usr/bin/python
	"""
	train_distillation.py
	Knowledge distillation: DLoc teacher -> lightweight student.

	Loss: L = ALPHA * MSE(student, ground_truth) + BETA * MSE(student, teacher_output)
	Default: ALPHA=0.7, BETA=0.3 (from professor's spec)

	Requires pre-computed teacher outputs (run precompute_teacher.py first).
	Set MODE='standalone' to train without distillation (ground truth only).

	Usage:
	python precompute_teacher.py # run once
	python train_distillation.py # distillation (default: mobilenet)

	To change student or mode, edit CONFIG below.
	"""
	import sys
	import os
	import time
	import torch
	import torch.nn as nn
	import numpy as np
	import warnings
	with warnings.catch_warnings():
	warnings.filterwarnings("ignore", category=FutureWarning)

	from utils import localization_error, write_log
	from data_loader import load_data
	from student_mobilenet import MobileNetStudent
	from student_tinycnn import TinyCNNStudent

	# ============================================================
	# CONFIG — edit these before running
	# ============================================================
	STUDENT_TYPE = 'tinycnn'
	MODE = 'distill'

	ALPHA = 0.7 # ground truth loss weight
	BETA = 0.3 # distillation loss weight (ignored in standalone)
	EPOCHS = 50
	BATCH_SIZE = 32
	LR = 1e-4
	NUM_WORKERS = 4

	TEACHER_TRAIN_PT = './data/teacher_train_outputs.pt'
	TEACHER_TEST_PT = './data/teacher_test_outputs.pt'

	# ============================================================
	# DATA PATHS (Fig 10b — Jacobs Hall)
	# ============================================================
	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',
	]

	# ============================================================
	# SETUP
	# ============================================================
	torch.manual_seed(0)
	np.random.seed(0)

	device = torch.device('cuda:0')
	save_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
	run_label = f'{MODE}_{STUDENT_TYPE}'
	save_dir = os.path.join('./runs', f'{run_label}_{save_name}')
	os.makedirs(save_dir, exist_ok=True)

	print(f'Mode: {MODE}')
	print(f'Student: {STUDENT_TYPE}')
	if MODE == 'distill':
	print(f'Loss: {ALPHA}L_GT + {BETA}L_KD')
	else:
	print(f'Loss: L_GT (standalone, no distillation)')
	print(f'Save dir: {save_dir}')

	# ============================================================
	# LOAD DATA INTO RAM
	# ============================================================
	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_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_w = torch.cat(test_w, dim=0)
	test_lbl = torch.cat(test_lbl, dim=0)
	print(f'Testing samples: {test_w.shape[0]}')

	# Load pre-computed teacher outputs if distilling
	if MODE == 'distill':
	print('Loading pre-computed teacher outputs...')
	train_teacher = torch.load(TEACHER_TRAIN_PT, weights_only=True)
	test_teacher = torch.load(TEACHER_TEST_PT, weights_only=True)
	assert train_teacher.shape[0] == train_w.shape[0], \
	f'Teacher/data mismatch: {train_teacher.shape[0]} vs {train_w.shape[0]}'
	print(f'Teacher outputs loaded: train={train_teacher.shape}, test={test_teacher.shape}')
	train_data = torch.utils.data.TensorDataset(train_w, train_lbl, train_teacher)
	test_data = torch.utils.data.TensorDataset(test_w, test_lbl, test_teacher)
	else:
	train_data = torch.utils.data.TensorDataset(train_w, train_lbl)
	test_data = torch.utils.data.TensorDataset(test_w, test_lbl)

	train_loader = torch.utils.data.DataLoader(
	train_data, batch_size=BATCH_SIZE, shuffle=True,
	num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True,
	)
	test_loader = torch.utils.data.DataLoader(
	test_data, batch_size=BATCH_SIZE, shuffle=False,
	num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True,
	)

	# ============================================================
	# CREATE STUDENT
	# ============================================================
	if STUDENT_TYPE == 'mobilenet':
	student = MobileNetStudent(input_nc=4).to(device)
	elif STUDENT_TYPE == 'tinycnn':
	student = TinyCNNStudent(input_nc=4).to(device)
	else:
	raise ValueError(f'Unknown student type: {STUDENT_TYPE}')

	n_params = sum(p.numel() for p in student.parameters())
	print(f'Student params: {n_params:,} ({n_params/1e6:.2f}M)')

	optimizer = torch.optim.Adam(student.parameters(), lr=LR)
	criterion = nn.MSELoss()

	# ============================================================
	# TRAINING LOOP
	# ============================================================
	print(f'Starting {MODE} training...')
	best_median = float('inf')

	for epoch in range(1, EPOCHS + 1):
	epoch_start = time.time()

	# --- TRAIN ---
	student.train()
	train_loss_sum = 0
	train_errors = []

	for data in train_loader:
	inputs = data[0].to(device)
	labels = data[1].to(device)

	student_output = student(inputs)
	L_GT = criterion(student_output, labels)

	if MODE == 'distill':
	teacher_out = data[2].to(device)
	L_KD = criterion(student_output, teacher_out)
	loss = ALPHA * L_GT + BETA * L_KD
	else:
	loss = L_GT

	optimizer.zero_grad()
	loss.backward()
	optimizer.step()

	train_loss_sum += loss.item()
	train_errors.extend(localization_error(
	student_output.detach().cpu().numpy(),
	labels.cpu().numpy(),
	scale=0.1,
	))

	train_loss_avg = train_loss_sum / len(train_loader)
	train_median = np.median(train_errors)

	# --- TEST ---
	student.eval()
	test_loss_sum = 0
	test_errors = []

	with torch.no_grad():
	for data in test_loader:
	inputs = data[0].to(device)
	labels = data[1].to(device)

	student_output = student(inputs)
	loss = criterion(student_output, labels)

	test_loss_sum += loss.item()
	test_errors.extend(localization_error(
	student_output.cpu().numpy(),
	labels.cpu().numpy(),
	scale=0.1,
	))

	test_loss_avg = test_loss_sum / len(test_loader)
	test_median = np.median(test_errors)
	test_p90 = np.percentile(test_errors, 90)
	test_p99 = np.percentile(test_errors, 99)
	elapsed = time.time() - epoch_start

	# --- LOG ---
	write_log([str(train_loss_avg)], run_label, log_dir=save_dir, log_type='train_loss')
	write_log([str(train_median)], run_label, log_dir=save_dir, log_type='train_median_error')
	write_log([str(test_loss_avg)], run_label, log_dir=save_dir, log_type='test_loss')
	write_log([str(test_median)], run_label, log_dir=save_dir, log_type='test_median_error')
	write_log([str(test_p90)], run_label, log_dir=save_dir, log_type='test_90_error')
	write_log([str(test_p99)], run_label, log_dir=save_dir, log_type='test_99_error')

	print(f'Epoch {epoch}/{EPOCHS} \| {elapsed:.0f}s \| '
	f'train_loss={train_loss_avg:.6f} train_med={train_median:.3f}m \| '
	f'test_med={test_median:.3f}m P90={test_p90:.3f}m P99={test_p99:.3f}m')

	# --- SAVE ---
	if test_median < best_median:
	best_median = test_median
	torch.save(student.state_dict(), os.path.join(save_dir, f'best_{STUDENT_TYPE}.pth'))
	print(f' -> New best: {best_median:.4f}m')

	torch.save(student.state_dict(), os.path.join(save_dir, f'latest_{STUDENT_TYPE}.pth'))

	print(f'\nDone. Best median error: {best_median:.4f}m')
	print(f'Results saved to: {save_dir}')