zc/mlp-project
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

...

Browse source
This commit is contained in:
Zhengyi Chen 2024-03-06 03:26:28 +00:00
parent ae9bc34fde
commit 0d35d607fe
2 changed files with 36 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

2
preprocess_data.py
View file

@ -52,7 +52,7 @@ def pre_dataset_sh():
# np.random.seed(0)
# random.seed(0)
for _, img_path in tqdm(img_paths, desc="Preprocessing Data"):
for img_path in tqdm(img_paths, desc="Preprocessing Data"):
img_data = cv2.imread(img_path)
mat = io.loadmat(
img_path

20
train.py
View file

@ -28,9 +28,9 @@ logger = logging.getLogger("train")
if not args.export_to_h5:
writer = SummaryWriter(args.save_path + "/tensorboard-run")
else:
train_df = pd.DataFrame(columns=["l1loss", "composite-loss"])
train_df = pd.DataFrame(columns=["l1loss", "composite-loss"], dtype=float)
train_stat_file = args.save_path + "/train_stats.h5"
test_df = pd.DataFrame(columns=["mse", "mae"])
test_df = pd.DataFrame(columns=["mse", "mae"], dtype=float)
test_stat_file = args.save_path + "/test_stats.h5"
@ -268,13 +268,13 @@ def train_one_epoch(
device_type = "cuda"
# Desperate measure to reduce mem footprint...
with torch.autocast(device_type):
# with torch.autocast(device_type):
# fpass
out, gt_count = model(img, kpoint)
# loss
loss = criterion(out, gt_count) # wrt. transformer
if args.export_to_h5:
train_df.loc[epoch * i, "l1loss"] = loss.item()
train_df.loc[epoch * i, "l1loss"] = float(loss.item())
else:
writer.add_scalar(
"L1-loss wrt. xformer (train)", loss, epoch * i
@ -291,7 +291,7 @@ def train_one_epoch(
)
)
if args.export_to_h5:
train_df.loc[epoch * i, "composite-loss"] = loss.item()
train_df.loc[epoch * i, "composite-loss"] = float(loss.item())
else:
writer.add_scalar("Composite loss (train)", loss, epoch * i)
@ -364,6 +364,9 @@ def valid_one_epoch(test_loader, model, device, epoch, args):
torch.sum(pred_count).item()
))
if args.debug:
break
mae = mae * 1.0 / (len(test_loader) * batch_size)
mse = np.sqrt(mse / (len(test_loader)) * batch_size)
if args.export_to_h5:
@ -373,14 +376,17 @@ def valid_one_epoch(test_loader, model, device, epoch, args):
else:
writer.add_scalar("MAE (valid)", mae, epoch)
writer.add_scalar("MSE (valid)", mse, epoch)
if len(xformed) != 0:
if len(xformed) != 0 and not args.export_to_h5:
img_grid = torchvision.utils.make_grid(xformed)
writer.add_image("STN: transformed image", img_grid, epoch)
if not args.export_to_h5:
writer.flush()
nni.report_intermediate_result(mae)
print("* MAE {mae:.3f} | MSE {mse:.3f} *".format(
mae=mae, mse=mse
))
writer.flush()
return mae