def eval_net(net, dataset, dir_img, dir_mask, args):
net.eval()
if args.gpu:
net.to(args.device)
total = 0
val = get_imgs_and_masks(dataset['val'], dir_img, dir_mask)
for i, b in enumerate(val):
img = np.array(b[0]).astype(np.float32)
mask = np.array(b[1]).astype(np.float32)
img = torch.from_numpy(img)[None, None, :, :]
mask = torch.from_numpy(mask).unsqueeze(0)
if args.gpu:
img = img.to(args.device)
mask = mask.to(args.device)
mask_pred = net(img)
mask_pred = (mask_pred > 0.5).float() # 得到预测的分割图
total += dice_coeff(mask_pred, mask, args.device).cpu().item()
current_score = total / (i + 1)
global best_score
print('current score is %f' % current_score)
print('best score is %f' % best_score)
if current_score > best_score:
best_score = current_score
print('current best score is {}'.format(best_score))
if args.save_cp:
print('saving checkpoint')
mkdir_p('checkpoint')
torch.save(net.state_dict(), './checkpoint/unet.pth')
return best_score
def get_result(net, gpu=False):
ids = get_ids(dir_img)
val = get_imgs_and_masks(ids, dir_img, dir_mask, 1.0)
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
def finetune(net,
optimizer,
criterion,
trainset,
log,
path,
iters=100,
epochs=None,
batch_size=2,
gpu=True,
scale=0.5):
net.train()
bce_meter = AverageMeter()
dir_img = 'data/train/'
dir_mask = 'data/train_masks/'
if epochs is None: # Fine-tune using iterations of mini-batches
epochs = 1
else: # Fine-tune using entire epochs
iters = None
for e in range(epochs):
# reset the generators
train = get_imgs_and_masks(trainset, dir_img, dir_mask, scale)
with tqdm(total=len(trainset)) as progress_bar:
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs).squeeze()
loss = criterion(masks_pred, true_masks)
bce_meter.update(loss.item(), batch_size)
optimizer.zero_grad()
loss.backward()
optimizer.step()
progress_bar.update(batch_size)
progress_bar.set_postfix(epoch=e, BCE=bce_meter.avg)
if i == 0 and e == 0:
log.info("FLOPs after pruning: \n{}".format(
flops_count(net, imgs.shape[2:])))
if i == iters: # Stop finetuning after sufficient mini-batches
break
log.info("Finished finetuning")
log.info("Finetuned loss: {}".format(bce_meter.avg))
torch.save(net.state_dict(), path)
log.info('Saving finetuned to {}...'.format(path))
def train(net, iddataset, dir_img, dir_mask, optimizer, criterion, args):
n_train = len(iddataset['train'])
# criterion = nn.BCELoss()
net.train()
if args.gpu:
net.to(args.device)
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask)
epoch_loss = 0.
for i, b in enumerate(batch(train, args.batchsize)):
imgs = np.array([data[0] for data in b]).astype(np.float32)
masks = np.array([data[1] for data in b]).astype(np.float32)
imgs = torch.from_numpy(imgs).unsqueeze(1).float()
masks = torch.from_numpy(masks).float()
if args.gpu:
imgs = imgs.to(args.device)
masks = masks.to(args.device)
optimizer.zero_grad()
mask_pred = net(imgs)
mask_prob_flat = mask_pred.view(-1)
masks_flat = masks.view(-1)
loss = criterion(mask_prob_flat, masks_flat)
epoch_loss += loss.cpu().item()
loss.backward()
optimizer.step()
print('training progress:{0:.4f} --- loss: {1:.6f}'.format(
i * args.batchsize / n_train, loss.item()))
def train_net(net,
device,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.15,
save_cp=True,
img_scale=0.5):
ids = get_ids(dir_img)
iddataset = split_train_val(ids, val_percent)
logging.info('''Starting training:
Epochs: {epochs}
Batch size: {batch_size}
Learning rate: {lr}
Training size: {len(iddataset["train"])}
Validation size: {len(iddataset["val"])}
Checkpoints: {save_cp}
Device: {device.type}
Images scaling: {img_scale}
''')
n_train = len(iddataset['train'])
n_val = len(iddataset['val'])
optimizer = optim.Adam(net.parameters(), lr=lr)
if net.n_classes > 1:
criterion = nn.CrossEntropyLoss()
else:
criterion = nn.BCEWithLogitsLoss()
for epoch in range(epochs):
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
f1_score = 0
num = 0
with tqdm(total=n_train, desc='Epoch {epoch + 1}/{epochs}',
unit='img') as pbar:
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
imgs = imgs.to(device=device)
true_masks = true_masks.to(device=device)
masks_pred = net(imgs)
# print('mask:',masks_pred.size())
# print('lab:',true_masks.size())
loss = criterion(masks_pred, true_masks)
masks_pred_np = masks_pred.detach().cpu().numpy()
true_masks_np = true_masks.detach().cpu().numpy()
epoch_loss += loss.item()
# print("----------------------------------------")
# print('masks_pred',type(masks_pred),masks_pred,'\n')
# print('true_masks',type(true_masks),true_masks,'\n')
# print('mask:',masks_pred.size(),'\n')
# print('lab:',true_masks.size(),'\n')
pre_2D = np.array(masks_pred_np[0][0])
true_2D = np.array(true_masks_np[0][0])
pre_2D_threhold = pre_2D
pre_2D_threhold[pre_2D_threhold > 0.5] = 1
pre_2D_threhold[pre_2D_threhold <= 0.5] = 0
# print("pre_2D.shape",pre_2D.shape,'\n')
# print("true_2D.shape" ,true_2D.shape,'\n')
# print("true_2D.flatten()",true_2D.flatten(),'\n')
# print("pre_2D.flatten()",pre_2D.flatten(),'\n')
pixel_accuracy = (pre_2D, true_2D)
f1_score += metrics.f1_score(true_2D.flatten(),
pre_2D_threhold.flatten())
num = num + 1
# print("----------------------------------------")
# val_score1 = eval_net1(net,val,device,n_val)
pbar.set_postfix(**{'loss (batch)': loss.item()})
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.update(batch_size)
if save_cp:
try:
os.mkdir(dir_checkpoint)
logging.info('Created checkpoint directory')
except OSError:
pass
torch.save(net.state_dict(),
dir_checkpoint + 'CP_epoch{epoch + 1}.pth')
logging.info('Checkpoint {epoch + 1} saved !')
val_score = eval_net(net, val, device, n_val)
f1_score /= num
print("f1-score:", f1_score, '\n')
if net.n_classes > 1:
logging.info('Validation cross entropy: {}'.format(val_score))
else:
logging.info('Validation Dice Coeff: {}'.format(val_score))
def train_net(net,
device,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.15,
save_cp=True,
img_scale=0.5):
ids = get_ids(dir_img)
iddataset = split_train_val(ids, val_percent)
n_train = len(iddataset['train'])
n_val = len(iddataset['val'])
optimizer = optim.Adam(net.parameters(), lr=lr)
if net.n_classes > 1:
criterion = nn.CrossEntropyLoss()
else:
criterion = nn.BCEWithLogitsLoss()
for epoch in range(epochs):
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
with tqdm(total=n_train,
desc='Epoch {0}/{1}'.format(epoch + 1, epochs),
unit='img') as pbar:
for i, b in enumerate(batch(train, batch_size)):
current_lr = adjust_learning_rate(optimizer, epoch, epochs,
pbar.n, n_train)
random_rate = 0
if epoch > epochs / 2:
random_rate = (epoch * 0.1) / epochs
b = custom_transforms.random_data_augmentation(
b, random_rate=random_rate)
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1][0] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
imgs = imgs.to(device=device)
true_masks = true_masks.to(device=device)
masks_pred = net(imgs)
loss = criterion(masks_pred, true_masks.long())
epoch_loss += loss.item()
pbar.set_postfix(
**{
'lr:{0}, random_rate:{1}, loss:'.format(
current_lr, random_rate):
loss.item()
})
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.update(batch_size)
if save_cp:
try:
os.mkdir(dir_checkpoint)
logging.info('Created checkpoint directory')
except OSError:
pass
torch.save(net.state_dict(),
dir_checkpoint + 'CP_epoch{0}.pth'.format(epoch + 1))
logging.info('Checkpoint {0} saved !'.format(epoch + 1))
val_score = eval_net(net, val, device, n_val)
if net.n_classes > 1:
logging.info('Validation cross entropy: {0}'.format(val_score))
else:
logging.info('Validation Dice Coeff: {0}'.format(val_score))
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = 'data/train/'
dir_mask = 'data/train_masks/'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
best_dice = 0.0
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
true_masks = true_masks / true_masks.max()
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs = F.sigmoid(masks_pred)
masks_probs_flat = masks_probs.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
with open('acc.log', 'a+') as w:
w.write('eopch {}, acc:{}\n'.format(epoch, val_dice))
if save_cp and val_dice > best_dice:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
best_dice = val_dice
from unet import UNet
import torch
ori_w, ori_h = 852, 480
dir_img = '/home/zhuzhu/Desktop/mid project/raw_data'
dir_mask = '/home/zhuzhu/Desktop/mid project/groundtruth'
ids = get_ids(dir_img)
iddataset = split_train_val(ids, 0.05)
net = UNet(1, 2)
net.eval()
net.load_state_dict(
torch.load(
'/media/zhuzhu/0C5809B80C5809B8/draft/unet/checkpoint/unet_0.854608765.pth',
map_location='cpu'))
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask)
c = 0
for i, b in enumerate(val):
img = np.array(b[0]).astype(np.float32)
mask = np.array(b[1]).astype(np.float32)
with torch.no_grad():
img = torch.from_numpy(img)[None, None, :, :]
mask = torch.from_numpy(mask).unsqueeze(0)
mask_pred = net(img)
coal, gangue = mask_pred.data.numpy().reshape(2, -1)
coal_fpr, coal_tpr, coal_th = roc_curve(
mask.squeeze().numpy()[0].reshape(-1), coal)
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
dir_img = '../all/reduced/'
dir_mask = '../all/masks/'
dir_checkpoint = 'checkpoints/'
# fetches ids in directory, without .jpg extension
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
# create a generator to try: yeild each value enumerate(batch(train, batch_size))
# or transfer enumerate() to try loop somehow.
epoch_loss = 0
exception_amount = 0
for i, b in enumerate(batch(train, batch_size)):
try:
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs = F.sigmoid(masks_pred)
masks_probs_flat = masks_probs.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(
i * batch_size / N_train, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
except KeyboardInterrupt:
torch.save(net.state_dict(), 'INTERRUPTED.pth')
print('Saved interrupt')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
except Exception as ah:
print(ah)
exception_amount += 1
if exception_amount > 20:
print('things arent going good')
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
device,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.15,
save_cp=True,
img_scale=0.5):
ids = get_ids(dir_img)
iddataset = split_train_val(ids, val_percent)
logging.info('''Starting training:
Epochs: {epochs}
Batch size: {batch_size}
Learning rate: {lr}
Training size: {len(iddataset["train"])}
Validation size: {len(iddataset["val"])}
Checkpoints: {save_cp}
Device: {device.type}
Images scaling: {img_scale}
''')
n_train = len(iddataset['train'])
n_val = len(iddataset['val'])
optimizer = optim.Adam(net.parameters(), lr=lr)
#optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.75)
criterion = nn.BCELoss()
for epoch in range(epochs):
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
with tqdm(total=n_train, desc='Epoch {epoch + 1}/{epochs}',
unit='img') as pbar:
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
imgs = imgs.to(device=device)
true_masks = true_masks.to(device=device)
masks_pred = net(imgs)
loss = criterion(masks_pred, true_masks)
epoch_loss += loss.item()
pbar.set_postfix(**{'loss (batch)': loss.item()})
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.update(batch_size)
if save_cp:
try:
os.mkdir(dir_checkpoint)
logging.info('Created checkpoint directory')
except OSError:
pass
torch.save(net.state_dict(), dir_checkpoint + 'gpu_3.pth')
logging.info('gpu_3 saved !')
val_dice = eval_net(net, val, device, n_val)
logging.info('Validation Dice Coeff: {}'.format(val_dice))
def train_net(net,
train_dir=None,
groundtruth_dir=None,
weight_dir=None,
weight_name='DeepInsthink',
val_percent=0.05,
epochs=5,
batch_size=1,
lr=0.1,
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = train_dir
dir_mask = groundtruth_dir
dir_checkpoint = weight_dir
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
max_DSC = 0
max_ep_checkpoint = 0
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
batchN = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('ep: {3:.0f} batch [{0:.0f}/{1:.0f}] - loss: {2:.6f}'.format(
i + 1, N_train / batch_size, loss.item(), epoch + 1))
#if (i % 5==0):
# val_dice = eval_net(net, val, gpu)
# print('Validation Dice Coeff: {}'.format(val_dice))
val_dice = eval_net(net, val, gpu)
print('Epoch {0:} -- Loss: {1:} -- Validation DSC: {2:}'.format(
epoch, epoch_loss / i, val_dice))
if (val_dice >= max_DSC):
max_DSC = val_dice
max_ep_checkpoint = epoch + 1
if save_cp:
torch.save(
net.state_dict(),
dir_checkpoint + weight_name + '-{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
print('Maximum checkpoint is ' + weight_name + '-{0:}' +
'with {1:} DSC'.format(max_ep_checkpoint, max_DSC))
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.2,
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = '/media/workspace/DATASET/Eyes/STARE/stare-images'
dir_mask = '/media/workspace/DATASET/Eyes/STARE/labels'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
ids = split_ids(ids, 4)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([w[0] for w in b]).astype(np.float32)
true_masks = np.array([w[1] for w in b])
true_masks[true_masks == 255] = 1
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
if not os.path.exists(dir_checkpoint):
os.mkdir(dir_checkpoint)
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_height=512,
img_scale=0.5):
#dir_img = 'carvana-image-masking-challenge/train/'
#dir_mask = 'carvana-image-masking-challenge/train_masks/'
dir_img = '/root/ctw/train_images_preprocess_other/'
dir_mask = '/root/ctw/train_images_mask_preprocess_other/'
#dir_img = '/root/ctw/val_images_preprocess_test/'
#dir_mask = '/root/ctw/val_images_mask_preprocess_test/'
dir_checkpoint = 'checkpoints/'
ids = list(get_ids(dir_img))
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.Adam(net.parameters(),lr=lr)
# optimizer = optim.SGD(net.parameters(),
# lr=lr,
# momentum=0.92,
# weight_decay=0.0005)
criterion = nn.BCELoss()
#criterion = nn.MSELoss()
#import scipy.misc
iteration = 0
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_height, img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_height, img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
#print(i, len(b))
"""
for j in b:
#print(j[0].shape, j[1].shape)
#print(j[1])
#scipy.misc.toimage(j[0], cmin=0.0, cmax=1.0).save('%s_outfile.jpg'%count)
#scipy.misc.toimage(j[1], cmin=0.0, cmax=1.0).save('%s_outmask.jpg'%count)
count += 1
"""
iteration += 1
try:
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
# print("\nImgs : \n{}".format(np.unique(imgs)))
# print("\ntrue mask \n {} ".format(np.unique(true_masks)))
#print('%s'%(datetime.datetime.now()), '{0:.4f}'.format(i * batch_size))
except Exception as e:
print(Exception)
continue
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
if iteration%100==0:
print('iter %s'%iteration, '%s'%(datetime.datetime.now()), '{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
val_iou=val_dice/(2-val_dice)
print('Validation Dice Coeff: {}'.format(val_dice))
print('Validation iouScore : {}'.format(val_iou))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def fit(net,
tf_writer,
epochs=5,
batch_size=1,
lr=0.0001,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=1,
l2=1e-8,
mom=0.9,
n_classes=4,
loss_function='bce',
alpha_non_zero=1,
resize_in=500):
# dir_png = "data/caddata_line_v2_1_mini/png"
# dir_mask = "data/caddata_line_v2_1_mini/mask"
dir_png_train = os.path.join(args.data, 'train', 'png')
dir_mask_train = os.path.join(args.data, 'train', 'mask')
dir_png_val = os.path.join(args.data, 'val', 'png')
dir_mask_val = os.path.join(args.data, 'val', 'mask')
# dir_mask = "data/mini/mask"
dir_checkpoint = os.path.join(args.log, 'checkpoints')
if not os.path.isdir(dir_checkpoint):
os.mkdir(dir_checkpoint)
# train
ids_train = get_ids(dir_png_train)
ids_train = split_ids(ids_train, n=1)
l_ids_train = list(ids_train)
# val
ids_val = get_ids(dir_png_val)
ids_val = split_ids(ids_val, n=1)
l_ids_val = list(ids_val)
# iddataset = split_train_val(ids, val_percent)
iddataset = {'train': l_ids_train, 'val': l_ids_val}
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
if loss_function == 'bce':
criterion = nn.BCELoss()
elif loss_function == 'mse':
criterion = nn.MSELoss()
# criterion = nn.CrossEntropyLoss()
with open(os.path.join(args.log, 'log.txt'), 'w+') as f_log:
for epoch in range(epochs):
# print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
# optimizer = optim.SGD(net.parameters(),
# lr=lr * (args.lr_decay ** epoch),
# momentum=mom,
# weight_decay=l2)
optimizer = optim.Adam(net.parameters(),
lr=lr * (args.lr_decay**epoch),
weight_decay=l2)
# optimizer = optim.RMSprop(net.parameters(),
# lr=lr * (args.lr_decay ** epoch),
# weight_decay=l2)
print("Current lr = {}".format(lr * (args.lr_decay**epoch)))
f_log.write("Current lr = {}".format(lr * (args.lr_decay**epoch)) +
'\n')
net.train()
# shuffle training set
random.shuffle(l_ids_train)
iddataset = {'train': l_ids_train, 'val': l_ids_val}
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_png_train,
dir_mask_train, img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_png_val,
dir_mask_val, img_scale)
epoch_loss = 0
epoch_tot = 0
epoch_acc = 0
epoch_acc_all = 0
i_out = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([j[0] for j in b]).astype(np.uint8)
true_masks = np.array([j[1] for j in b]).astype(np.uint8)
imgs_2 = np.zeros(
(imgs.shape[0], N_CHANNELS, resize_in, resize_in))
true_masks_2 = np.zeros(
(true_masks.shape[0], resize_in, resize_in, n_classes))
## data augmentation
for j in range(imgs.shape[0]):
img = imgs[j, :, :, :]
mask = true_masks[j, :, :, :]
# print(np.unique(mask))
pil_img = Image.fromarray(img, 'RGB')
pil_mask = Image.fromarray(mask, 'RGB')
##debug##
# pil_img.show()
# pil_mask.show()
# print(np.unique(mask))
seed = np.random.randint(124521346)
# Resize Crop ratio: img
random.seed(seed)
pil_img = torchvision.transforms.RandomResizedCrop(
size=(resize_in),
scale=(0.8, 1.0),
interpolation=Image.NEAREST)(pil_img)
# Resize Crop ratio: true_masks
random.seed(seed)
pil_mask = torchvision.transforms.RandomResizedCrop(
size=(resize_in),
scale=(0.8, 1.0),
interpolation=Image.NEAREST)(pil_mask)
"""
# rotate seed
random_degree = randrange(360)
# rotate: img
pil_img = torchvision.transforms.functional.rotate(pil_img, angle=random_degree)
# rotate: true_masks
pil_mask = torchvision.transforms.functional.rotate(pil_mask, angle=random_degree)
"""
# color: img
# color: true_masks
arr_img = rgb_pil_to_bw_norm_arr(pil_img)
# print(np.unique(arr_img))
imgs_2[j, N_CHANNELS - 1, :, :] = arr_img
# print(np.unique(imgs_2))
arr_mask = np.array(pil_mask)
arr_mask_2 = np.zeros((resize_in, resize_in, n_classes))
for in_c in range(n_classes):
in_1, in_2 = np.where(arr_mask[:, :, 0] == (in_c + 1))
arr_mask_2[in_1, in_2,
in_c] = arr_mask[in_1, in_2,
0].astype(bool).astype(
np.float32)
true_masks_2[j, :, :, :] = arr_mask_2
# true_masks_2.astype(np.float32)
# print("======")
# print(np.unique(true_masks_2))
## To TorchTensor
# imgs:
# imgs = torch.from_numpy(imgs_2.astype(np.float32).transpose(0,3,1,2))
imgs = torch.from_numpy(imgs_2.astype(np.float32))
# true_masks:
# true_masks = torch.from_numpy(true_masks_2.transpose(0,3,1,2))
true_masks = torch.from_numpy(true_masks_2)
# imgs = torch.from_numpy(imgs)
# true_masks = torch.from_numpy(true_masks)
# true_masks = np.transpose(true_masks, (0, 3, 1, 2))
# assert imgs.size()[1] == N_CHANNELS
# assert true_masks.size()[1] == n_classes
# assert true_masks.size()[2] == imgs.size()[2]
# assert true_masks.size()[3] == imgs.size()[3]
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
# view(-1)
masks_probs_flat = masks_pred.view(-1)
if gpu:
true_masks_flat = true_masks.view(-1)
else:
true_masks_flat = true_masks.reshape(-1)
true_masks_flat = true_masks_flat.float()
# loss
loss = criterion(masks_probs_flat, true_masks_flat)
in_nonzero = torch.nonzero(true_masks_flat)
loss_nonzero = criterion(masks_probs_flat[in_nonzero],
true_masks_flat[in_nonzero])
if in_nonzero.size(0) != 0:
loss = loss + alpha_non_zero * loss_nonzero
epoch_loss += loss.item()
true_masks_flat_bin = true_masks_flat.unsqueeze(0)
masks_probs_flat_bin = (masks_probs_flat >
0.5).float().unsqueeze(0)
this_dice = dice_coeff(masks_probs_flat_bin,
true_masks_flat_bin).item()
epoch_tot += this_dice
# e = np.array(masks_probs_flat_bin.cpu())
# f = np.array(true_masks_flat_bin.cpu())
acc_train = iou(np.array(true_masks_flat_bin.cpu()),
np.array(masks_probs_flat_bin.cpu()))
acc_train_all = iou_all(np.array(true_masks_flat_bin.cpu()),
np.array(masks_probs_flat_bin.cpu()))
epoch_acc += acc_train
epoch_acc_all += acc_train_all
if i % print_interval == print_interval - 1:
print(
'{0} / {1} steps. --- loss: {2:.6f}, IoU_train_nz: {3:.4f}, IoU_train_all: {4:.4f}, dice: {5:.4f}'
.format(i, int(N_train / batch_size),
epoch_loss / (i + 1), epoch_acc / (i + 1),
epoch_acc_all / (i + 1), epoch_tot / (i + 1)))
f_log.write(
'{0} / {1} steps. --- loss: {2:.6f}, ACC_train: {3:.4f}, IoU_train_all: {4:.4f}, dice: {4:.4f}'
.format(i, int(N_train / batch_size), epoch_loss /
(i + 1), epoch_acc / (i + 1), epoch_acc_all /
(i + 1), epoch_tot / (i + 1)) + '\n')
optimizer.zero_grad()
loss.backward()
optimizer.step()
i_out = i
print(
'Epoch {} finished ! Loss: {}, IoU: {}, IoU_all: {}, dice: {}'.
format(epoch, epoch_loss / (i_out + 1),
epoch_acc / (i_out + 1), epoch_acc_all / (i_out + 1),
epoch_tot / (i + 1)))
f_log.write(
'Epoch finished ! Loss: {}, IoU: {}, IoU_all: {}, dice: {}'.
format(epoch_loss / (i_out + 1), epoch_acc /
(i_out + 1), epoch_acc_all / (i_out + 1), epoch_tot /
(i + 1)) + '\n')
tf_writer.add_scalar('data/train_loss', epoch_loss / (i_out + 1),
epoch)
tf_writer.add_scalar('data/train_iou', epoch_acc / (i_out + 1),
epoch)
tf_writer.add_scalar('data/train_iou_all',
epoch_acc_all / (i_out + 1), epoch)
tf_writer.add_scalar('data/train_dice', epoch_tot / (i_out + 1),
epoch)
## Evaluate
net.eval()
tot_val = 0
epoch_loss_val = 0
epoch_acc_val = 0
epoch_acc_val_all = 0
for i_val, b_val in enumerate(batch(val, batch_size)):
# imgs_val = np.array([j[0] for j in b_val]).astype(np.float32)
# true_masks_val = np.array([j[1] for j in b_val])
#
# imgs_val = torch.from_numpy(imgs_val)
# true_masks_val = torch.from_numpy(true_masks_val)
# true_masks_val = np.transpose(true_masks_val, (0, 3, 1, 2))
####
imgs_val = np.array([j[0] for j in b_val]).astype(np.uint8)
true_masks_val = np.array([j[1]
for j in b_val]).astype(np.uint8)
imgs_2 = np.zeros(
(imgs_val.shape[0], N_CHANNELS, resize_in, resize_in))
true_masks_2 = np.zeros(
(true_masks_val.shape[0], resize_in, resize_in, n_classes))
## data augmentation
if imgs_val.shape[0] == batch_size:
for j in range(imgs_val.shape[0]):
img = imgs_val[j, :, :, :]
mask = true_masks_val[j, :, :, :]
# print(np.unique(mask))
pil_img = Image.fromarray(img, 'RGB')
pil_mask = Image.fromarray(mask, 'RGB')
# resize
pil_img = torchvision.transforms.Resize(
size=(resize_in),
interpolation=Image.NEAREST)(pil_img)
pil_mask = torchvision.transforms.Resize(
size=(resize_in),
interpolation=Image.NEAREST)(pil_mask)
# upload img and mask to imgs_2 and mask_2
arr_img = rgb_pil_to_bw_norm_arr(pil_img)
imgs_2[j, N_CHANNELS - 1, :, :] = arr_img
arr_mask = np.array(pil_mask)
arr_mask_2 = np.zeros(
(resize_in, resize_in, n_classes))
for in_c in range(n_classes):
in_1, in_2 = np.where(arr_mask[:, :,
0] == (in_c + 1))
arr_mask_2[in_1, in_2,
in_c] = arr_mask[in_1, in_2,
0].astype(bool).astype(
np.float32)
true_masks_2[j, :, :, :] = arr_mask_2
# true_masks_2.astype(np.float32)
# print("======")
# print(np.unique(true_masks_2))
## To TorchTensor
# imgs:
imgs_val = torch.from_numpy(imgs_2.astype(np.float32))
# true_masks:
true_masks_val = torch.from_numpy(true_masks_2)
if gpu:
imgs_val = imgs_val.cuda()
true_masks_val = true_masks_val.cuda()
masks_pred_val = net(imgs_val)
masks_probs_flat_val = masks_pred_val.view(-1)
if gpu:
true_masks_flat_val = true_masks_val.view(-1)
else:
true_masks_flat_val = true_masks_val.reshape(-1)
true_masks_flat_val = true_masks_flat_val.float()
true_masks_flat_bin_val = true_masks_flat_val.unsqueeze(0)
masks_probs_flat_bin_val = (masks_probs_flat_val >
0.5).float().unsqueeze(0)
dice_val = dice_coeff(masks_probs_flat_bin_val,
true_masks_flat_bin_val).item()
acc_val = iou(np.array(true_masks_flat_bin_val.cpu()),
np.array(masks_probs_flat_bin_val.cpu()))
acc_val_all = iou_all(
np.array(true_masks_flat_bin_val.cpu()),
np.array(masks_probs_flat_bin_val.cpu()))
epoch_acc_val += acc_val
epoch_acc_val_all += acc_val_all
tot_val += dice_val
loss_val = criterion(masks_probs_flat_val,
true_masks_flat_val)
in_nonzero = torch.nonzero(true_masks_flat_val)
loss_val_nonzero = criterion(
masks_probs_flat_val[in_nonzero],
true_masks_flat_val[in_nonzero])
if in_nonzero.size(0) != 0:
loss_val = loss_val + alpha_non_zero * loss_val_nonzero
epoch_loss_val = loss_val / (i_val + 1)
epoch_dice_val = tot_val / (i_val + 1)
epoch_acc_val = epoch_acc_val / (i_val + 1)
epoch_acc_val_all = epoch_acc_val_all / (i_val + 1)
# val_dice = eval_net(net, val, gpu)
print(
'* Val: Loss: {0:.6f}, IoU: {1:.3f}, IoU_all: {2:.3f}, Dice: {3:.3f}'
.format(epoch_loss_val, epoch_acc_val, epoch_acc_val_all,
epoch_dice_val))
f_log.write(
'* Val: Loss: {0:.6f}, IoU: {1:.3f}, IoU_all: {2:.3f}, Dice: {3:.3f}'
.format(epoch_loss_val, epoch_acc_val, epoch_acc_val_all,
epoch_dice_val) + '\n')
tf_writer.add_scalar('data/val_loss', epoch_loss_val, epoch)
tf_writer.add_scalar('data/val_iou', epoch_acc_val, epoch)
tf_writer.add_scalar('data/val_iou_all', epoch_acc_val_all, epoch)
tf_writer.add_scalar('data/val_dice', epoch_dice_val, epoch)
if save_cp and (epoch % save_interval == save_interval - 1):
torch.save(net.state_dict(),
dir_checkpoint + '/CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
f_log.write('Checkpoint {} saved !'.format(epoch + 1) + '\n')
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.5,
save_cp=True,
gpu=True,
img_scale=0.5):
dir_img = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_few_images/'
dir_mask = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_masks_few_images/'
# dir_img = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_images/'
# dir_mask = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_masks/'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
#optimizer = optim.SGD(net.parameters(),
# lr=lr,
# momentum=0.9,
# weight_decay=0.0005)
optimizer = optim.Adam(net.parameters(),
lr = lr,
betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch {} finished ! Loss: {}'.format(epoch, epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_prune_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5,
num_prune_iterations=100):
dir_img = 'data/train/'
dir_mask = 'data/train_masks/'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.eval()
sub = list(enumerate(val))[:3]
for i, b in sub:
img = b[0]
true_mask = b[1]
img = torch.from_numpy(img).unsqueeze(0)
true_mask = torch.from_numpy(true_mask).unsqueeze(0)
if gpu:
img = img.cuda()
true_mask = true_mask.cuda()
mask_pred = net(img)[0]
mask_pred = (mask_pred > 0.5).float()
val_dice = dice_coeff(mask_pred, true_mask).item()
print('Validation Dice Coeff at batch {}: {}'.format(
i + 1, val_dice))
net.train()
for i, b in enumerate(batch(train, batch_size)):
if i > 2:
break
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
for j in range(num_prune_iterations):
if j == num_prune_iterations - 1:
net.prune(verbose=True)
else:
net.prune(verbose=False)
summary(net, (3, 640, 640))
for i, b in sub:
img = b[0]
true_mask = b[1]
img = torch.from_numpy(img).unsqueeze(0)
true_mask = torch.from_numpy(true_mask).unsqueeze(0)
if gpu:
img = img.cuda()
true_mask = true_mask.cuda()
mask_pred = net(img)[0]
mask_pred = (mask_pred > 0.5).float()
val_dice = dice_coeff(mask_pred, true_mask).item()
print('Validation Dice Coeff at batch {}: {}'.format(
i + 1, val_dice))
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
# val_dice = eval_net(net, val, gpu)
# print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
epochs=30,
batch_size=6,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
# dir_img = 'E:/A_paper_thesis/paper5/tensorflow_deeplabv3plus_scrapingData/dataset/Scraping_Data2/train_db'
dir_img = 'data/train_db/'
dir_mask = 'data/GT_bw/'
dir_checkpoint = 'checkpoint0919/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale)
epoch_loss = 0
epoch_iou = 0
epoch_xor=0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('step:', i)
# print('Validation Dice Coeff: {0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))
print('Validation Dice Coeff: {0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))
# # mean iou
# intersect = sum(masks_probs_flat*true_masks_flat)
# union = sum(masks_probs_flat+true_masks_flat)
# iou = (intersect+0.001)/(union-intersect+0.001)
# epoch_iou +=iou
# mean iou
smooth = 1e-6 # we smooth to avoid our devision 0/0
intersect = sum(masks_probs_flat*true_masks_flat)
union = sum(masks_probs_flat+true_masks_flat)-intersect
iou = (intersect+smooth)/(union+smooth)
epoch_iou +=iou
# calculate xor
# xor quation is: xor = (union(output hợp ground truth) - intersect(output giao ground truth))/ ground truth
# xor = (union-intersect)/ground truth
xor = (union - intersect)/sum(true_masks_flat)
epoch_xor += xor
print('mean IoU: {:.4f}'.format(iou))
# print('mean IoU1: {:.4f}'.format(iou1))
print('mean xor: {:.4f}'.format(xor))
# end of mean iou
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! epoch_Loss: {:.6f}'.format(epoch_loss / i))
print('epoch_iou: {:.4f}'.format(epoch_iou / i))
print('epoch_xor: {:.4f}'.format(epoch_xor / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('epoch_Validation Dice Coeff: {:.4f}'.format(val_dice))
# need to write mean iou of evaluate here(reference val_dice)
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = 'data/train/' # 训练图像文件夹
dir_mask = 'data/train_masks/' # 图像的结果文件夹
dir_checkpoint = 'checkpoints/' # 训练好的网络保存文件夹
ids = get_ids(dir_img) # 图片名字的后4位为数字，能作为图片id
# 得到元祖列表为[(id1,0),(id1,1),(id2,0),(id2,1),...,(idn,0),(idn,1)]
# 这样的作用是后面重新设置生成器时会通过后面的0,1作为utils.py中get_square函数的pos参数，pos=0的取左边的部分，pos=1的取右边的部分
# 这样图片的数量就会变成2倍
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent) # 将数据分为训练集和验证集两份
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train']) # 训练集长度
optimizer = optim.SGD(
net.parameters(), # 定义优化器
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss() # 损失函数
for epoch in range(epochs): # 开始训练
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train() # 设置为训练模式
# reset the generators重新设置生成器
# 对输入图片dir_img和结果图片dir_mask进行相同的图片处理，即缩小、裁剪、转置、归一化后，将两个结合在一起，返回(imgs_normalized, masks)
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32) # 得到输入图像数据
true_masks = np.array([i[1] for i in b]) # 得到图像结果数据
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs) # 图像输入的网络后得到结果masks_pred，结果为灰度图像
masks_probs_flat = masks_pred.view(-1) # 将结果压扁
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat) # 对两个结果计算损失
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss /
i)) # 一次迭代后得到的平均损失
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(args, net, val_percent=0.05, save_cp=True):
dir_img = os.path.join(args.dataset_folder, 'data/train/')
dir_mask = os.path.join(args.dataset_folder, 'data/train_masks/')
dir_checkpoint = os.path.join(args.dataset_folder, 'checkpoints/')
if not os.path.exists(dir_checkpoint):
os.makedirs(dir_checkpoint)
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
'''.format(args.epochs, args.batch_size, args.lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(args.epochs):
print('Starting epoch {}/{}.'.format(args.epochs + 1, args.epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
args.img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
args.img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, args.batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
# if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(
i * args.batch_size / N_train, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
epochs=5,
batch_size=1,
lr=1e-3,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = '/home/xyj/data/spacenet/vegas/images_rgb_1300/'
dir_mask = '/home/xyj/test/Pytorch-UNet/data/train_mask_point/'
dir_checkpoint = 'checkpoints_point/'
if not os.path.exists(dir_checkpoint):
os.mkdir(dir_checkpoint)
# ids = get_ids(dir_img) # 返回train文件夹下文件的名字列表，生成器（except last 4 character，.jpg这样的）
with open('train_list.txt', 'r') as f:
lines = f.readlines()
ids = (i.strip('\n')[:-4] for i in lines)
ids = split_ids(
ids) # 返回(id, i), id属于ids，i属于range(n)，相当于把train的图✖️了n倍多张，是tuple的生成器
iddataset = split_train_val(
ids, val_percent
) # validation percentage，是dict = {"train": ___(一个list), "val"：___(一个list)}
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
# optimizer = optim.SGD(net.parameters(),
# lr=lr,
# momentum=0.9,
# weight_decay=0.0005)
optimizer = optim.Adam(net.parameters(),
lr=lr,
betas=(0.9, 0.999),
eps=1e-3)
# scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=40,gamma = 0.3)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] // 200 for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
# scheduler.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
epochs=20,
batch_size=1,
lr=0.1,
lrd=0.99,
val_percent=0.05,
save_cp=True,
gpu=True,
img_scale=0.5,
imagepath='',
maskpath='',
cpsavepath=''):
dir_img = imagepath
dir_mask = maskpath
dir_checkpoint = cpsavepath
classweight = [1, 2, 3, 2]
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
logname = cpsavepath + '/' + 'losslog.txt'
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
# classweight = [1,4,8,4]
criterion = BCELoss_weight(classweight)
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
with open(logname, "a") as f:
f.write('Starting epoch {}/{}.'.format(epoch + 1, epochs) + "\n")
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
lr = lr * lrd
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('lr', lr)
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
true_masks = np.transpose(true_masks, axes=[0, 3, 1, 2])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
# print('masks_pred.shape',masks_pred.shape)
# print('true_masks.shape', true_masks.shape)
masks_probs_flat = masks_pred
true_masks_flat = true_masks
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
printinfo = '{0:.4f} --- loss: {1:.6f}'.format(
i * batch_size / N_train, loss.item())
print(printinfo)
with open(logname, "a") as f:
f.write(printinfo + "\n")
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
with open(logname, "a") as f:
f.write('Epoch finished ! Loss: {}'.format(epoch_loss / i) + "\n")
if 1:
val_dice = eval_net(net, val)
print('Validation Dice Coeff: {}'.format(val_dice))
with open(logname, "a") as f:
f.write('Validation Dice Coeff: {}'.format(val_dice) + "\n")
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
with open(logname, "a") as f:
f.write('Checkpoint {} saved !'.format(epoch + 1) + "\n")
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5,
dir_img=None,
dir_mask=None,
dir_checkpoint=None,
channels=1,
classes=1):
ids = os.listdir(dir_img)
if not os.path.exists(dir_checkpoint):
os.makedirs(dir_checkpoint, mode=0o755)
iddataset = split_train_val(ids, val_percent)
print('Starting training:')
print('Epochs: ' + str(epochs))
print('Batch size: ' + str(batch_size))
print('Learning rate: ' + str(lr))
print('Training size: ' + str(len(iddataset['train'])))
print('Validation size: ' + str(len(iddataset['val'])))
print('Checkpoints: ' + str(save_cp))
N_train = len(iddataset['train'])
optimizer = optim.RMSprop(net.parameters(), lr=lr)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
# Run Batch
for i, b in enumerate(batch(train, batch_size)):
# Grab data
try:
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
except:
print(
'prob have dimension issues, wrong orientations or half reconned images'
)
# Deal with dimension issues
if channels == 1:
imgs = np.expand_dims(imgs, 1)
if classes > 1:
true_masks = to_categorical(true_masks, num_classes=classes)
# Play in torch's sandbox
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
# Send to GPU
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
# Predicted segmentations
masks_pred = net(imgs)
# Flatten
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
# Calculate losses btwn true/predicted
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
# Batch Loss
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
# Backprop
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Epoch Loss
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, epoch, dir_checkpoint, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(
net.state_dict(),
os.path.join(dir_checkpoint, 'CP{}.pth'.format(epoch + 1)))
print('Checkpoint {} saved !'.format(epoch + 1))
from utils import split_train_val, batch, get_imgs_and_masks, inputImageViz
# get all values that you require from your input
args = get_args()
dir_img = args.images
dir_mask = args.masks
dir_checkpoint = 'checkpoints/'
val_percent = args.valpercent
scale = args.scale
maxWidth = args.maxWidth
maxHeight = args.maxHeight
n_classes= args.classes
iddataset = split_train_val(dir_img, val_percent)
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, scale, maxWidth, maxHeight, n_classes)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, scale, maxWidth, maxHeight, n_classes)
for epoch in range(args.epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
#net.train()
for i, k in enumerate(batch(val, args.batchsize)):
print(i)
imgs = np.array([i[0] for i in k]).astype(np.float32)
true_masks = np.array([i[1] for i in k])
# comment this if required
inputImageViz (imgs, true_masks)
#------------ your training code here!!!--------------------
def train_net(args,
net,
epochs,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
global best_dice, best_loss
dir_img = '/home/mori/Programming/Net_Pruning/unetdataset_patchImg/img/'
dir_mask = '/home/mori/Programming/Net_Pruning/unetdataset_patchImg/graylabel/'
ids = get_ids(dir_img) # get file name (without .png)
print("ids:{}".format(ids))
ids = split_ids(ids) # 重采样？
print("ids:{}".format(ids))
iddataset = split_train_val(ids, val_percent) # 按给定比例划分打乱的数据集
###### count parameters ############
paras = sum([p.data.nelement() for p in net.parameters()])
print('''
Starting training:
Epochs: {}
Parameters: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
Deepsupervision: {}
'''.format(epochs, paras, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu),
str(args.deepsupervision)))
N_train = len(iddataset['train'])
print("N_train:{}".format(N_train))
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
New_lr = adjust_learning_rate(optimizer, epoch, epochs)
print(' lr: {}'.format(New_lr))
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
best_iou = 0
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)): # 手动分出batch
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
true_masks_flat = true_masks.view(-1)
true_masks_flat = true_masks_flat / 255 # 归一化
output = net(imgs)
masks_pred = F.sigmoid(output)
if args.deepsupervision:
#### unet++ with deepsupervision
loss = 0
for mp in masks_pred:
masks_probs_flat = mp.view(-1)
loss += criterion(masks_probs_flat, true_masks_flat)
loss /= len(masks_pred)
epoch_loss += loss.item()
else:
masks_probs_flat = masks_pred.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
## todo: adjust iou
iou = iou_score(output, true_masks / 255)
######## record the best iou
if iou > best_iou:
best_iou = iou
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
newloss = loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
print('Best iou: {}'.format(best_iou))
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
writer.add_scalar('train_loss', epoch_loss / i, (epoch + 1))
writer.add_scalar('val_dice', val_dice, (epoch + 1))
writer.add_scalar('best iou', best_iou, (epoch + 1))
if save_cp:
#torch.save(net.state_dict(),dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
#print('Checkpoint {} saved !'.format(epoch + 1))
dice_best = val_dice > best_dice
loss_best = epoch_loss / i < best_loss
best_dice = max(val_dice, best_dice)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': net.state_dict(),
'best_dice': best_dice,
'best_loss': best_loss,
}, dice_best, loss_best)
print('Best dice: ', best_dice)
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.2,
save_cp=True,
gpu=False,
img_scale=0.5):
path = [['data/ori1/', 'data/gt1/'],
['data/original1/', 'data/ground_truth1/'],
['data/Original/', 'data/Ground_Truth/']]
dir_img = path[0][0]
dir_mask = path[0][1]
dir_checkpoint = 'sdgcheck/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.7,
weight_decay=0.005)
'''
optimizer = optim.Adam(net.parameters(),
lr=lr,
weight_decay=0.0005)
'''
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
x = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
'''
ori=np.transpose(imgs[0], axes=[1, 2, 0])
scipy.misc.imsave("ori/ori_"+str(x)+'.jpg', ori)
gt = np.stack((true_masks[0],)*3, axis=-1)
#gt=np.transpose(true_masks[0], axes=[1, 2, 0])
scipy.misc.imsave("gt/gt_"+str(x)+'.jpg', gt)
'''
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
x += 1
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(
net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05, # 训练集：验证集= 0.95： 0.05
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = opt_train.dir_img
dir_mask = opt_train.dir_mask
dir_checkpoint = opt_train.dir_checkpoint
# 得到 图片路径列表 ids为 图片名称（无后缀名）
ids = get_ids(dir_img)
# 得到truple元组 （无后缀名的 图片名称，序号）
# eg:当n为2 图片名称为bobo.jpg 时, 得到（bobo,0） （bobo,1）
# 当序号为0 时，裁剪宽度，得到左边部分图片 当序号为1 时，裁剪宽度，得到右边部分图片
ids = split_ids(ids)
# 打乱数据集后，按照val_percent的比例来 切分 训练集 和 验证集
iddataset = split_train_val(ids, val_percent)
print('''
开始训练:
Epochs: {}
Batch size: {}
Learning rate: {}
训练集大小: {}
验证集大小: {}
GPU: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(gpu)))
#训练集大小
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
#二进制交叉熵
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
# reset the generators
# 每轮epoch得到 训练集 和 验证集
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
# 重置epoch损失计数器
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
# 得到 一个batch的 imgs tensor 及 对应真实mask值
# 当序号为0 时，裁剪宽度，得到左边部分图片[384,384,3] 当序号为1 时，裁剪宽度，得到右边部分图片[384,190,3]
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
# 将值转为 torch tensor
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
# 训练数据转到GPU上
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
# 得到 网络输出的预测mask [10,1,384,384]
masks_pred = net(imgs)
# 经过sigmoid
masks_probs = F.sigmoid(masks_pred)
masks_probs_flat = masks_probs.view(-1)
true_masks_flat = true_masks.view(-1)
# 计算二进制交叉熵损失
loss = criterion(masks_probs_flat, true_masks_flat)
# 统计一个epoch的所有batch的loss之和，用以计算 一个epoch的 loss均值
epoch_loss += loss.item()
# 输出 当前epoch的第几个batch 及 当前batch的loss
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
# 优化器梯度清零
optimizer.zero_grad()
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 一轮epoch结束，该轮epoch的 loss均值
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
# 每轮epoch之后使用验证集进行评价
if True:
# 评价函数：Dice系数 Dice距离用于度量两个集合的相似性
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
# 保存模型
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
indent=4,
sort_keys=True)))
# Dataset
if not os.path.exists(splitfile): # Our constant datasplit
ids = get_ids(dir_img) # [file1, file2]
ids = split_ids(ids) # [(file1, 0), (file1, 1), (file2, 0), ...]
iddataset = split_train_val(ids, 0.2, splitfile)
log.info("New split dataset")
else:
with open(splitfile) as f:
iddataset = json.load(f)
log.info("Load split dataset")
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
args.scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, args.scale)
# Model Initialization
net = UNet(n_channels=3, n_classes=1, f_channels=args.channel_txt)
log.info("Built model using {}...".format(args.channel_txt))
if args.gpu:
net.cuda()
if args.load:
net.load_state_dict(torch.load(args.load))
log.info('Loading checkpoint from {}...'.format(args.load))
pruner = Pruner(net) # Pruning handler
criterion = nn.BCELoss()
# Ranking on the train dataset
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
dir_img = '../dataset/train/images/'
dir_mask = '../dataset/train/masks/'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print(('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu))))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
print(('Starting epoch {}/{}.'.format(epoch + 1, epochs)))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)): # b[batch_id][0/1]: a batch of image(0)+mask(1)
#print(('b[0]',b[0][0].shape,b[0][1].shape))
#imgs = []
#for img_msk in b:
# imgs.append(img_msk[0])
#print(len(imgs))
#imgs = np.array(imgs) # Wrong: not all images are of the same shape
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print(('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(('Epoch finished ! Loss: {}'.format(epoch_loss / i)))
if 1:
val_dice = eval_net(net, val, gpu)
print(('Validation Dice Coeff: {}'.format(val_dice)))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'cropped_CP{}.pth'.format(epoch + 1))
print(('Checkpoint {} saved !'.format(epoch + 1)))
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=0.5):
# dir_img = 'data/train/'
# dir_mask = 'data/train_masks/'
dir_img = 'E:/git/dataset/tgs-salt-identification-challenge/train/images/'
dir_mask = 'E:/git/dataset/tgs-salt-identification-challenge/train/masks/'
# dir_img = 'E:/git/dataset/tgs-salt-identification-challenge/train/my_images/'
# dir_mask = 'E:/git/dataset/tgs-salt-identification-challenge/train/my_masks/'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
# true_masks = np.array([i[1] for i in b])#np.rot90(m)
true_masks = np.array([i[1].T / 65535 for i in b]) #np.rot90(m)
# show_batch_image(true_masks)
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
# show_batch_image(imgs)
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train,
loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net,
epochs=50,
batch_size=1,
lr=0.1,
val_percent=0.1,
save_cp=True,
gpu=True,
img_scale=[513, 513]):
#dir_img = '/home/lixiaoxing/data/DRIVE/train/'
#dir_mask = '/home/lixiaoxing/data/DRIVE/trainannot/'
dir_img = '/home/lixiaoxing/github/Pytorch-UNet/data/DRIVE/AV_groundTruth/training/images_jpg/'
dir_mask = '/home/lixiaoxing/github/Pytorch-UNet/data/DRIVE/AV_groundTruth/training/vessel/'
dir_checkpoint = 'checkpoints/'
if os.path.exists(dir_checkpoint) is False:
os.makedirs(dir_checkpoint)
ids = get_ids(dir_img)
ids = split_ids(ids)
iddataset = split_train_val(ids, val_percent)
print('''
Starting training:
Epochs: {}
Batch size: {}
Learning rate: {}
Training size: {}
Validation size: {}
Checkpoints: {}
CUDA: {}
'''.format(epochs, batch_size, lr, len(iddataset['train']),
len(iddataset['val']), str(save_cp), str(gpu)))
N_train = len(iddataset['train'])
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=0.0005)
criterion = nn.BCELoss()
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
net.train()
# reset the generators
train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask,
img_scale)
val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask,
img_scale)
#print(train)
epoch_loss = 0
for i, b in enumerate(batch(train, batch_size)):
imgs = np.array([i[0] for i in b]).astype(np.float32)
true_masks = np.array([i[1] for i in b])
imgs = torch.from_numpy(imgs)
true_masks = torch.from_numpy(true_masks)
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
# up features
#print('**********************up**************************')
#up_feature = net.extract_features(imgs)
#print(up_feature.shape)
#ff = net._blocks[38]._depthwise_conv
masks_pred = net(imgs)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)
#print(true_masks_flat.shape)
#print(masks_probs_flat.shape)
loss = criterion(masks_probs_flat, true_masks_flat)
epoch_loss += loss.item()
#print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / i))
if 1:
val_dice = eval_net(net, val, gpu)
print('Validation Dice Coeff: {}'.format(val_dice))
if save_cp:
torch.save(net.state_dict(),
dir_checkpoint + 'CP{}.pth'.format(epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))