def CE_Net_Train(train_i=0):
NAME = 'fold' + str(i + 1) + '_6CE-Net' + Constants.ROOT.split('/')[-1]
solver = MyFrame(CE_Net_, dice_bce_loss, 2e-4)
batchsize = torch.cuda.device_count() * Constants.BATCHSIZE_PER_CARD #4
# For different 2D medical image segmentation tasks, please specify the dataset which you use
# for examples: you could specify "dataset = 'DRIVE' " for retinal vessel detection.
txt_train = 'fold' + str(train_i + 1) + '_train.csv'
txt_test = 'fold' + str(train_i + 1) + '_test.csv'
dataset_train = MyDataset(txt_path=txt_train,
transform=transforms.ToTensor(),
target_transform=transforms.ToTensor())
dataset_test = MyDataset(txt_path=txt_test,
transform=transforms.ToTensor(),
target_transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(dataset,
batchsize=batchsize,
shuffle=True,
num_workers=2)
test_loader = torch.utils.data.DataLoader(dataset,
batchsize=batchsize,
shuffle=False,
num_workers=2)
# start the logging files
mylog = open('logs/' + NAME + '.log', 'w')
no_optim = 0
total_epoch = Constants.TOTAL_EPOCH # 300
train_epoch_best_loss = Constants.INITAL_EPOCH_LOSS # 10000
best_test_score = 0
for epoch in range(1, total_epoch + 1):
data_loader_iter = iter(train_loader)
data_loader_test = iter(test_loader)
train_epoch_loss = 0
index = 0
tic = time()
# train
for img, mask in data_loader_iter:
solver.set_input(img, mask)
train_loss, pred = solver.optimize()
train_epoch_loss += train_loss
index = index + 1
# test
test_sen = 0
test_ppv = 0
test_score = 0
for img, mask in data_loader_test:
solver.set_input(img, mask)
pre_mask, _ = solver.test_batch()
test_score += dice_coeff(y_test, pre_mask, False)
test_sen += sensitive(y_test, pre_mask)
# test_sen = test_sen.cpu().data.numpy()
test_ppv += positivepv(y_test, pre_mask)
# test_ppv = test_ppv.cpu().data.numpy()
print(test_sen / len(data_loader_test),
test_ppv / len(data_loader_test),
test_score / len(data_loader_test))
# solver.set_input(x_test, y_test)
# pre_mask, _ = solver.test_batch()
# test_score = dice_coeff(y_test, pre_mask, False)
# test_sen = sensitive(y_test, pre_mask)
# test_sen = test_sen.cpu().data.numpy()
# test_ppv = positivepv(y_test, pre_mask)
# test_ppv = test_ppv.cpu().data.numpy()
# print('111111111111111111111',type(test_score))
# # show the original images, predication and ground truth on the visdom.
# show_image = (img + 1.6) / 3.2 * 255.
# viz.img(name='images', img_=show_image[0, :, :, :])
# viz.img(name='labels', img_=mask[0, :, :, :])
# viz.img(name='prediction', img_=pred[0, :, :, :])
if test_score > best_test_score:
print('1. the dice score up to ', test_score, 'from ',
best_test_score, 'saving the model')
best_test_score = test_score
solver.save('./weights/' + NAME + '.th')
train_epoch_loss = train_epoch_loss / len(data_loader_iter)
# print(mylog, '********')
print('epoch:',
epoch,
' time:',
int(time() - tic),
'train_loss:',
train_epoch_loss.cpu().data.numpy(),
file=mylog,
flush=True)
print('test_dice_loss: ',
test_score,
'test_sen: ',
test_sen,
'test_ppv: ',
test_ppv,
'best_score is ',
best_test_score,
file=mylog,
flush=True)
print('********')
print('epoch:', epoch, ' time:', int(time() - tic), 'train_loss:',
train_epoch_loss.cpu().data.numpy())
print('test_dice_score: ', test_score, 'test_sen: ', test_sen,
'test_ppv: ', test_ppv, 'best_score is ', best_test_score)
# print('train_loss:', train_epoch_loss)
# print('SHAPE:', Constants.Image_size)
if train_epoch_loss >= train_epoch_best_loss:
no_optim += 1
else:
no_optim = 0
train_epoch_best_loss = train_epoch_loss
# solver.save('./weights/' + NAME + '.th')
# if no_optim > Constants.NUM_EARLY_STOP:
# print(mylog, 'early stop at %d epoch' % epoch)
# print('early stop at %d epoch' % epoch)
# break
if no_optim > Constants.NUM_UPDATE_LR:
if solver.old_lr < 5e-7:
break
if solver.old_lr > 5e-4:
solver.load('./weights/' + NAME + '.th')
solver.update_lr(1.5, factor=True, mylog=mylog)
print('Finish!', file=mylog, flush=True)
print('Finish!')
mylog.close()
def train(train_i=0):
NAME = 'D2F5_fold'+str(train_i+1)+'_FPN.th'
print(NAME)
batchsize = 4
txt_train = 'N5fold'+str(train_i+1)+'_train.csv'
txt_test = 'N5fold'+str(train_i+1)+'_test.csv'
dataset_train = MyDataset(root='/home/wangke/ultrasound_data2/', txt_path=txt_train, transform=transforms.ToTensor(), target_transform=transforms.ToTensor())
dataset_test = MyDataset(root='/home/wangke/ultrasound_data2/', txt_path=txt_test, transform=transforms.ToTensor(), target_transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batchsize, shuffle=True, num_workers=2)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=batchsize, shuffle=False, num_workers=2, drop_last=True)
mylog = open('logs/'+NAME+'.log', 'w')
# model = FPN_Net(1, 1)
# summary(model)
slover = MyFrame(FPN_Net, dice_bce_loss, 2e-4)
total_epoch = 100
no_optim = 0
train_epoch_best_loss = 10000
best_test_score = 0
for epoch in range(1, total_epoch+1):
data_loader_iter = iter(train_loader)
data_loader_test = iter(test_loader)
train_epoch_loss = 0
index = 0
tic = time()
train_score = 0
for img, mask in data_loader_iter:
slover.set_input(img, mask)
train_loss, pred = slover.optimize()
train_score += dice_coeff(mask, pred.cpu().data, False)
train_epoch_loss +=train_loss
index +=1
test_sen = 0
test_ppv = 0
test_score = 0
test_acc = 0
test_spe = 0
test_f1s = 0
for img, mask in data_loader_test:
slover.set_input(img, mask)
pre_mask, _ = slover.test_batch()
test_score += dice_coeff(mask, pre_mask, False)
test_sen += sensitive(mask, pre_mask)
test_ppv += precision(mask, pre_mask)
test_acc += accuracy(mask, pre_mask)
test_spe += specificity(mask, pre_mask)
test_f1s += f1_score(mask, pre_mask)
test_sen /= len(data_loader_test)
test_ppv /= len(data_loader_test)
test_score /= len(data_loader_test)
test_acc /= len(data_loader_test)
test_spe /= len(data_loader_test)
test_f1s /= len(data_loader_test)
if test_score>best_test_score:
print('1. the dice score up to ', test_score, 'from ', best_test_score, 'saving the model', file=mylog, flush=True)
print('1. the dice score up to ', test_score, 'from ', best_test_score, 'saving the model')
best_test_score = test_score
slover.save('./weights/'+NAME+'.th')
train_epoch_loss = train_epoch_loss/len(data_loader_iter)
train_score = train_score/len(data_loader_iter)
print('epoch:', epoch, ' time:', int(time() - tic), 'train_loss:', train_epoch_loss.cpu().data.numpy(), 'train_score:', train_score, file=mylog, flush=True)
print('test_dice_loss: ', test_score, 'test_sen: ', test_sen.numpy(), 'test_ppv: ', test_ppv.numpy(), 'test_acc: ', test_acc.numpy(), 'test_spe: ', test_spe.numpy(), 'test_f1s: ', test_f1s.numpy(), 'best_score is ', best_test_score, file=mylog, flush=True)
print('********')
print('epoch:', epoch, ' time:', int(time() - tic), 'train_loss:', train_epoch_loss.cpu().data.numpy(), 'train_score:', train_score)
print('test_dice_loss: ', test_score, 'test_sen: ', test_sen.numpy(), 'test_ppv: ', test_ppv.numpy(), 'test_acc: ', test_acc.numpy(), 'test_spe: ', test_spe.numpy(), 'test_f1s: ', test_f1s.numpy(), 'best_score is ', best_test_score)
if train_epoch_loss >= train_epoch_best_loss:
no_optim +=1
else:
no_optim =0
train_epoch_best_loss = train_epoch_loss
print('Finish!', file=mylog, flush=True)
print('Finish!')
mylog.close()
def train(train_i=0):
NAME = 'fold' + str(train_i + 1) + '_deeplabv3_plus.th'
# slover = MyFrame(FSP_Net, dice_bce_loss, 2e-4)
slover = MyFrame(DeepLabV3Plus, dice_bce_loss, 5e-4)
batchsize = 4
txt_train = 'fold' + str(train_i + 1) + '_train.csv'
txt_test = 'fold' + str(train_i + 1) + '_test.csv'
dataset_train = MyDataset(txt_path=txt_train,
transform=transforms.ToTensor(),
target_transform=transforms.ToTensor())
dataset_test = MyDataset(txt_path=txt_test,
transform=transforms.ToTensor(),
target_transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(dataset_train,
batch_size=batchsize,
shuffle=True,
num_workers=2)
test_loader = torch.utils.data.DataLoader(dataset_test,
batch_size=batchsize,
shuffle=False,
num_workers=2)
mylog = open('logs/' + NAME + '.log', 'w')
total_epoch = 100
no_optim = 0
train_epoch_best_loss = 10000
best_test_score = 0
for epoch in range(1, total_epoch + 1):
data_loder_iter = iter(train_loader)
data_loder_test = iter(test_loader)
train_epoch_loss = 0
index = 0
tic = time()
train_score = 0
for img, mask in data_loder_iter:
slover.set_input(img, mask)
train_loss, pred = slover.optimize()
train_score += dice_coeff(mask, pred, False)
train_epoch_loss += train_loss
index += 1
test_sen = 0
test_ppv = 0
test_score = 0
for img, mask in data_loder_test:
slover.set_input(img, mask)
pre_mask, _ = slover.test_batch()
test_score += dice_coeff(mask, pre_mask, False)
test_sen += sensitive(mask, pre_mask)
test_ppv += positivepv(mask, pre_mask)
test_sen /= len(data_loder_test)
test_ppv /= len(data_loder_test)
test_score /= len(data_loder_test)
if test_score > best_test_score:
print('1. the dice score up to ',
test_score,
'from ',
best_test_score,
'saving the model',
file=mylog,
flush=True)
print('1. the dice score up to ', test_score, 'from ',
best_test_score, 'saving the model')
best_test_score = test_score
slover.save('./weights/' + NAME + '.th')
train_epoch_loss = train_epoch_loss / len(data_loder_iter)
train_score = train_score / len(data_loder_iter)
print('epoch:',
epoch,
' time:',
int(time() - tic),
'train_loss:',
train_epoch_loss.cpu().data.numpy(),
'train_score:',
train_score,
file=mylog,
flush=True)
print('test_dice_loss: ',
test_score,
'test_sen: ',
test_sen,
'test_ppv: ',
test_ppv,
'best_score is ',
best_test_score,
file=mylog,
flush=True)
print('epoch:', epoch, ' time:', int(time() - tic), 'train_loss:',
train_epoch_loss.cpu().data.numpy(), 'train_score:', train_score)
print('test_dice_loss: ',
test_score,
'test_sen: ',
test_sen,
'test_ppv: ',
test_ppv,
'best_score is ',
best_test_score,
file=mylog)
if train_epoch_loss >= train_epoch_best_loss:
no_optim += 1
else:
no_optim = 0
train_epoch_best_loss = train_epoch_loss
print('Finish!', file=mylog, flush=True)
print('Finish!')
mylog.close()
