Spaces:
Running
Running
| import torch | |
| from torch.utils.data import DataLoader | |
| import numpy as np | |
| import random | |
| from tqdm import tqdm | |
| import argparse | |
| from model import SegRoot | |
| from dataloader import StaticTrainDataset, TestDataset, TrainDataset, LoopSampler | |
| from utils import ( | |
| dice_score, | |
| init_weights, | |
| evaluate, | |
| get_ids, | |
| load_vgg16, | |
| set_random_seed, | |
| ) | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--seed", default=42, type=int, help="set random seed") | |
| parser.add_argument("--width", default=8, type=int, help="width of SegRoot") | |
| parser.add_argument("--depth", default=5, type=int, help="depth of SegRoot") | |
| parser.add_argument("--bs", default=64, type=int, help="batch size of dataloaders") | |
| parser.add_argument("--lr", default=1e-2, type=float, help="learning rate") | |
| parser.add_argument("--epochs", default=200, type=int, help="max epochs of training") | |
| parser.add_argument( | |
| "--verbose", default=5, type=int, help="intervals to save and validate model" | |
| ) | |
| parser.add_argument( | |
| "--dynamic", action="store_true", help="use dynamic sub-images during training" | |
| ) | |
| def train_one_epoch(model, train_iter, optimizer, device): | |
| model.train() | |
| for p in model.parameters(): | |
| p.requires_grad = True | |
| for x, y in train_iter: | |
| x, y = x.to(device), y.to(device) | |
| bs = x.shape[0] | |
| optimizer.zero_grad() | |
| y_pred = model(x) | |
| loss = 1 - dice_score(y, y_pred) | |
| loss = torch.sum(loss) / bs | |
| loss.backward() | |
| optimizer.step() | |
| if __name__ == "__main__": | |
| args = parser.parse_args() | |
| seed = args.seed | |
| bs = args.bs | |
| lr = args.lr | |
| width = args.width | |
| depth = args.depth | |
| epochs = args.epochs | |
| verbose = args.verbose | |
| # set random seed | |
| set_random_seed(seed) | |
| # define the device for training | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| # get training ids | |
| train_ids, valid_ids, test_ids = get_ids(65) | |
| # define dataloaders | |
| if args.dynamic: | |
| train_data = TrainDataset(train_ids) | |
| train_iter = DataLoader( | |
| train_data, batch_size=bs, num_workers=6, sampler=LoopSampler | |
| ) | |
| else: | |
| train_data = StaticTrainDataset(train_ids) | |
| train_iter = DataLoader(train_data, batch_size=bs, num_workers=6, shuffle=True) | |
| train_tdata = TestDataset(train_ids) | |
| valid_tdata = TestDataset(valid_ids) | |
| test_tdata = TestDataset(test_ids) | |
| # define model | |
| model = SegRoot(width, depth).to(device) | |
| model = model.apply(init_weights) | |
| # define optimizer and lr_scheduler | |
| optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4) | |
| scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( | |
| optimizer, mode="max", factor=0.5, verbose=True, patience=5 | |
| ) | |
| print(f"Start training SegRoot-({width},{depth}))......") | |
| print(f"Random seed is {seed}, batch size is {bs}......") | |
| print(f"learning rate is {lr}, max epochs is {epochs}......") | |
| best_valid = float("-inf") | |
| for epoch in tqdm(range(epochs)): | |
| train_one_epoch(model, train_iter, optimizer, device) | |
| if epoch % verbose == 0: | |
| train_dice = evaluate(model, train_tdata, device) | |
| valid_dice = evaluate(model, valid_tdata, device) | |
| scheduler.step(valid_dice) | |
| print( | |
| "Epoch {:05d}, train dice: {:.4f}, valid dice: {:.4f}".format( | |
| epoch, train_dice, valid_dice | |
| ) | |
| ) | |
| if valid_dice > best_valid: | |
| best_valid = valid_dice | |
| test_dice = evaluate(model, test_tdata, device) | |
| print("New best validation, test dice: {:.4f}".format(test_dice)) | |
| torch.save( | |
| model.state_dict(), | |
| f"../weights/best_segroot-({args.width},{args.depth}).pt", | |
| ) | |