def train(epochs, batch_size, learning_rate):
train_loader = torch.utils.data.DataLoader(SegThorDataset(
"data",
phase='train',
transform=transforms.Compose([Rescale(0.25),
Normalize(),
ToTensor()]),
target_transform=transforms.Compose([Rescale(0.25),
ToTensor()])),
batch_size=batch_size,
shuffle=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet().to(device)
model.apply(weight_init)
#optimizer = optim.Adam(model.parameters(), lr=learning_rate) #learning rate to 0.001 for initial stage
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.95)
#optimizer = adabound.AdaBound(params = model.parameters(), lr = 0.001, final_lr = 0.1)
for epoch in range(epochs):
print('Epoch {}/{}'.format(epoch + 1, epochs))
print('-' * 10)
running_loss = 0.0
loss_seg = np.zeros(5)
for batch_idx, (train_data, labels) in enumerate(train_loader):
train_data, labels = train_data.to(
device, dtype=torch.float), labels.to(device,
dtype=torch.uint8)
print("train data size", train_data.size())
print("label size", labels.size())
optimizer.zero_grad()
output = model(train_data)
print("output: {} and taget: {}".format(output.size(),
labels.size()))
loss_label, loss = dice_loss(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
for i in range(4):
loss_seg[i] += loss_label[i]
print("Length: ", len(train_loader))
epoch_loss = running_loss / len(train_loader)
epoch_loss_class = np.true_divide(loss_seg, len(train_loader))
print(
"Dice per class: Background = {:.4f} Eusophagus = {:.4f} Heart = {:.4f} Trachea = {:.4f} Aorta = {:.4f}\n"
.format(epoch_loss_class[0], epoch_loss_class[1],
epoch_loss_class[2], epoch_loss_class[3],
epoch_loss_class[4]))
print("Total Dice Loss: {:.4f}\n".format(epoch_loss))
os.makedirs("models", exist_ok=True)
torch.save(model, "models/model.pt")
def define_G(input_nc,
output_nc,
ngf,
norm='batch',
use_dropout=False,
gpu_ids=[]):
netG = None
use_gpu = len(gpu_ids) > 0
norm_layer = get_norm_layer(norm_type=norm)
if use_gpu:
assert (torch.cuda.is_available())
#netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9, gpu_ids=gpu_ids)
#netG = GeneratorUNet(in_channels=2, out_channels=1).cuda()
netG = UNet(n_classes=1).cuda()
if len(gpu_ids) > 0:
netG.cuda(gpu_ids[0])
netG.apply(weights_init)
return netG
def main():
opt = parser.parse_args()
print(opt)
print(torch.__version__)
opt.seed = 1337
random.seed(opt.seed)
torch.manual_seed(opt.seed)
np.random.seed(opt.seed)
torch.cuda.manual_seed(opt.seed)
series = [
f for f in os.listdir(opt.train_path[0])
if os.path.isdir(os.path.join(opt.train_path[0], f))
]
series.sort()
print("===> Building model")
enc_layers = [1, 2, 2, 4]
dec_layers = [1, 1, 1, 1]
number_of_channels = [int(8 * 2**i) for i in range(1, 1 + len(enc_layers))]
model = UNet(depth=len(enc_layers),
encoder_layers=enc_layers,
decoder_layers=dec_layers,
number_of_channels=number_of_channels,
number_of_outputs=3)
model.apply(weight_init.weight_init)
model = torch.nn.DataParallel(module=model, device_ids=range(opt.gpus))
trainer = train.Trainer(model=model,
name=opt.name,
models_root=opt.models_path,
rewrite=False)
trainer.cuda()
gc.collect()
opt.seed = 1337 # random.randint(1, 10000)
random.seed(opt.seed)
torch.manual_seed(opt.seed)
np.random.seed(opt.seed)
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
print("===> Loading datasets")
print('Train data:', opt.train_path)
series_val = [
'BraTS19_2013_0_1', 'BraTS19_2013_12_1', 'BraTS19_2013_16_1',
'BraTS19_2013_2_1', 'BraTS19_2013_23_1', 'BraTS19_2013_26_1',
'BraTS19_2013_29_1', 'BraTS19_CBICA_AAB_1', 'BraTS19_CBICA_AAP_1',
'BraTS19_CBICA_AMH_1', 'BraTS19_CBICA_AQD_1', 'BraTS19_CBICA_ATX_1',
'BraTS19_CBICA_AZH_1', 'BraTS19_CBICA_BHB_1', 'BraTS19_TCIA12_101_1',
'BraTS19_TCIA01_150_1', 'BraTS19_TCIA10_152_1', 'BraTS19_TCIA04_192_1',
'BraTS19_TCIA08_205_1', 'BraTS19_TCIA06_211_1', 'BraTS19_TCIA02_222_1',
'BraTS19_TCIA12_298_1', 'BraTS19_TCIA13_623_1', 'BraTS19_CBICA_ANV_1',
'BraTS19_CBICA_BBG_1', 'BraTS19_TMC_15477_1'
]
series_train = [f for f in series if f not in series_val]
print('Train {}'.format(series_train))
print('Val {}'.format(series_val))
train_set = dataloader.SimpleReader(
paths=opt.train_path,
patch_size=(144, 144, 128),
series=[
series_train,
] + [None for i in range(len(opt.train_path) - 1)],
annotation_path=opt.annotation_path,
images_in_epoch=8000,
patches_from_single_image=1)
val_set = dataloader.FullReader(path=opt.train_path[0], series=series_val)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True,
drop_last=True,
worker_init_fn=worker_init_fn)
batch_sampler = Data.BatchSampler(sampler=Data.SequentialSampler(val_set),
batch_size=1,
drop_last=True)
evaluation_data_loader = DataLoader(dataset=val_set,
num_workers=0,
batch_sampler=batch_sampler)
criterion = [
loss.Dice_loss_joint(index=0, priority=1).cuda(),
loss.BCE_Loss(index=0, bg_weight=1e-2).cuda(),
]
print("===> Building model")
print("===> Training")
trainer.train(criterion=criterion,
optimizer=optim.Adam,
optimizer_params={"lr": 2e-5,
"weight_decay": 1e-6,
"amsgrad": True,
},
scheduler=torch.optim.lr_scheduler.StepLR,
scheduler_params={"step_size": 16000,
"gamma": 0.5,
},
training_data_loader=training_data_loader,
evaluation_data_loader=evaluation_data_loader,
split_into_tiles=False,
pretrained_weights=None,
train_metrics=[metrics.Dice(name='Dice', input_index=0, target_index=0, classes=4), \
],
val_metrics=[metrics.Dice(name='Dice', input_index=0, target_index=0, classes=4),
metrics.Hausdorff_ITK(name='Hausdorff_ITK', input_index=0, target_index=0, classes=4),
],
track_metric='Dice',
epoches=opt.nEpochs,
default_val=np.array([0, 0, 0, 0, 0]),
comparator=lambda x, y: np.min(x) + np.mean(x) > np.min(y) + np.mean(y),
eval_cpu=False,
continue_form_pretraining=False
)
def train(epochs, batch_size, learning_rate):
torch.manual_seed(1234)
train_loader = torch.utils.data.DataLoader(SegThorDataset(
"/home/WIN-UNI-DUE/smnemada/Master_Thesis/SegThor/data/train",
phase='train',
transform=transforms.Compose([Rescale(1.0),
Normalize(),
ToTensor()])),
batch_size=batch_size,
shuffle=True)
'''
# Loading validation data
val_set = SegThorDataset("/home/WIN-UNI-DUE/smnemada/Master_Thesis/SegThor/data_val", phase='val',
transform=transforms.Compose([
Rescale(0.5),
Normalize(),
ToTensor2()
]))
val_loader = torch.utils.data.DataLoader(dataset=val_set,
batch_size=1,
shuffle=False)
'''
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet().to(device)
model.apply(weight_init)
#optimizer = optim.Adam(model.parameters(), lr=learning_rate) #learning rate to 0.001 for initial stage
optimizer = optim.SGD(model.parameters(),
lr=0.01,
momentum=0.9,
weight_decay=0.00001)
#optimizer = adabound.AdaBound(params = model.parameters(), lr = 0.001, final_lr = 0.1)
for epoch in range(epochs):
f = open('train_output.log', 'a')
f.write('Epoch {}/{}\n'.format(epoch + 1, epochs))
f.write('-' * 10)
running_loss = 0.0
running_loss_label = np.zeros(5)
for batch_idx, sample in enumerate(train_loader):
train_data, labels = sample['image'].to(
device,
dtype=torch.float), sample['label'].to(device,
dtype=torch.uint8)
optimizer.zero_grad()
output = model(train_data)
loss_label, loss = dice_loss2(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
for i in range(5):
running_loss_label[i] += loss_label[i]
epoch_loss = running_loss / len(train_loader)
writer.add_scalar('Train/Loss', epoch_loss, epoch)
f.write("\n Total Dice Loss: {:.4f}\n".format(epoch_loss))
epoch_loss_class = np.true_divide(running_loss_label,
len(train_loader))
f.write(
"Dice per class: Background = {:.4f} Eusophagus = {:.4f} Heart = {:.4f} Trachea = {:.4f} Aorta = {:.4f}\n"
.format(epoch_loss_class[0], epoch_loss_class[1],
epoch_loss_class[2], epoch_loss_class[3],
epoch_loss_class[4]))
#f.write("Dice per class: Background = {:.4f} Eusophagus = {:.4f}\n".format(epoch_loss_class[0], epoch_loss_class[1]))
f.close()
if epoch % 4 == 0:
os.makedirs("models", exist_ok=True)
torch.save(model, "models/model.pt")
# export scalar data to JSON for external processing
writer.export_scalars_to_json("./all_scalars.json")
writer.close()
os.makedirs("models", exist_ok=True)
torch.save(model, "models/model.pt")
def main(epochs, batch_size, learning_rate):
total_train_loss = 0
total_val_loss = 0
train_loss_seg = np.zeros(5)
val_loss_seg = np.zeros(5)
for train_list, test_list in k_folds(n_splits=4, subjects=41):
# Loading train data
train_loader = torch.utils.data.DataLoader(SegThorDataset(
"/home/WIN-UNI-DUE/smnemada/Master_Thesis/SegThor/data_sub/train_cv",
phase='train',
transform=transforms.Compose(
[Rescale(1.0), Normalize(),
ToTensor2()]),
file_list=train_list),
batch_size=batch_size,
shuffle=False)
# Loading validation data
val_set = SegThorDataset(
"/home/WIN-UNI-DUE/smnemada/Master_Thesis/SegThor/data_sub/train_cv",
phase='val',
transform=transforms.Compose(
[Rescale(1.0), Normalize(),
ToTensor2()]),
file_list=test_list)
val_loader = torch.utils.data.DataLoader(dataset=val_set,
batch_size=1,
shuffle=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet().to(device)
model.apply(weight_init)
#optimizer = optim.Adam(model.parameters(), lr=learning_rate) #learning rate to 0.001 for initial stage
optimizer = optim.SGD(model.parameters(),
lr=0.01,
momentum=0.9,
weight_decay=0.00001)
#optimizer = adabound.AdaBound(params = model.parameters(), lr = 0.001, final_lr = 0.1)
for epoch in range(epochs):
print('Epoch {}/{}'.format(epoch + 1, epochs))
print('-' * 10)
train_loss, train_loss_label = train(train_loader, model,
optimizer, epoch, device,
batch_size, train_list)
val_loss, val_loss_label = validation(val_loader, model, epoch,
device, batch_size,
test_list)
if val_loss < train_loss:
os.makedirs("models", exist_ok=True)
torch.save(model, "models/model.pt")
# Save model output
#save_results(epoch, device)
if epoch % 4 == 0:
os.makedirs("models", exist_ok=True)
torch.save(model, "models/model.pt")
total_train_loss = total_train_loss + train_loss
total_val_loss = total_val_loss + val_loss
train_loss_seg = train_loss_seg + train_loss_label
val_loss_seg = val_loss_seg + val_loss_label
#evaluate_model(epoch, device)
train_loss_seg = np.true_divide(train_loss_seg, 4)
val_loss_seg = np.true_divide(val_loss_seg, 4)
print(" Training Loss: ")
print(total_train_loss // epochs)
print(
"Background = {:.4f} Eusophagus = {:.4f} Heart = {:.4f} Trachea = {:.4f} Aorta = {:.4f}\n"
.format(train_loss_seg[0], train_loss_seg[1], train_loss_seg[2],
train_loss_seg[3], train_loss_seg[4]))
print(" Validation Loss: ")
print(total_val_loss // epochs)
print(
"Background = {:.4f} Eusophagus = {:.4f} Heart = {:.4f} Trachea = {:.4f} Aorta = {:.4f}\n"
.format(val_loss_seg[0], val_loss_seg[1], val_loss_seg[2],
val_loss_seg[3], val_loss_seg[4]))
