def train():
graphs, labels = load_data("datasets/train")
train_inputs, train_targets, val_inputs, val_targets = utils.split_train_val(
graphs, labels, val_rate=0.3)
model = GNNModel(8)
loss_func = BinaryCrossEntropy()
optimizer = Adam()
batch_generator = utils.BatchGenerator(batch_size=32)
min_loss = 100000
for epoch in range(50):
print(f"Epoch{epoch + 1}")
train_losses = []
for inputs, targets in batch_generator.generator(
train_inputs, train_targets):
train_loss, loss_grad = loss_func(model,
inputs,
targets,
is_grad=True)
optimizer.update(model, loss_grad)
train_losses.append(train_loss)
train_mean_loss = np.mean(train_losses)
pred = np.array([model.predict(input_)
for input_ in train_inputs]).squeeze()
train_accuracy = accuracy(pred, train_targets)
val_losses = []
for inputs, targets in batch_generator.generator(
val_inputs, val_targets):
val_loss, _ = loss_func(model, inputs, targets, is_grad=False)
val_losses.append(val_loss)
val_mean_loss = np.mean(val_losses)
pred = np.array([model.predict(input_)
for input_ in val_inputs]).squeeze()
val_accuracy = accuracy(pred, val_targets)
if min(min_loss, val_mean_loss) < min_loss:
min_loss = val_mean_loss
print(
f"Train loss: {train_mean_loss}\tTrain accuracy: {train_accuracy}"
)
print(
f"Validation loss: {val_mean_loss}\tValidation accuracy: {val_accuracy}"
)
print("")
def init(config):
''' Loads the GloVe embeddings for the words
which occur in the IMDB train set vocab
and uses that vocab to create train, validation
and test sets for the IMDB dataset. Extracts the
pad_id token.
'''
import os
if not os.path.isdir('.data'):
os.mkdir('.data')
# Extract the initial vocab from the IMDB dataset
vocab = IMDB(data_select='train')[0].get_vocab()
# Create GloVe embeddings based on original vocab
# word freqs
glove_vocab = torchtext.vocab.Vocab(
counter=vocab.freqs,
max_size=MAX_VOCAB_SIZE,
min_freq=MIN_FREQ,
vectors=torchtext.vocab.GloVe(name='6B'))
# Acquire 'Spacy' tokenizer for the vocab words
tokenizer = get_tokenizer('spacy', 'en_core_web_sm')
# Acquire train and test IMDB sets with previously created
# GloVe vocab and 'Spacy' tokenizer
train_set, test_set = IMDB(tokenizer=tokenizer, vocab=glove_vocab)
# Extract the vocab of the acquired train set
vocab = train_set.get_vocab()
# Extract the token used for padding
pad_id = vocab['<pad>']
# Split the train set into train and validation sets
train_set, valid_set = split_train_val(train_set)
config['train'] = train_set
config['val'] = valid_set
config['test'] = test_set
config['vocab'] = vocab
config['pad_id'] = pad_id
if __name__ == '__main__':
# 命令行参数
cfg_filename = sys.argv[1] # 配置文件名
mode = sys.argv[2] # 运行模式，包括train,test,online_test
# 读取配置文件
cfg = Config.fromfile(cfg_filename)
print("config filename: " + str(cfg_filename))
set_seed(cfg.random_seed)
# 划分数据集
if mode == 'split_data':
split_train_val()
# 训练模型
if mode == 'train':
train_main(cfg=cfg)
# 测试模型
if mode == 'test':
test_main(cfg=cfg)
# 复赛线上测试
if mode == 'test_online':
test_online_main(cfg=cfg)
print('end')
# plt.plot(hist_coal, 'r-')
# plt.plot(hist_gangue, 'g-')
# plt.xlim([0, 256])
# plt.show()
from sklearn.metrics import precision_recall_curve, roc_curve, classification_report
from utils import get_imgs_and_masks, get_ids, split_train_val
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():
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))
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))
for img in list(tif.iter_images()):
tiff = img
cv2.imwrite(os.path.join(dir_path, name[0] + '.jpg'))
img_path = os.path.join(dir_path, name[0] + '.jpg')
xml_path = os.path.join(dir_path, name[0] + '.xml')
img_table, fm_table = preprocess(img_path, xml_path)
assert len(img_table) == len(
fm_table), "image and feature map not match"
for idx in range(len(img_table)):
image_name.append(name[0] + '_' + str(idx))
cv2.imwrite(
os.path.join(save_path, name[0] + '_' + str(idx) + '.jpg'),
img_table[idx])
cv2.imwrite(
os.path.join(save_path,
name[0] + '_' + str(idx) + '_maskh.png'),
fm_table[idx][0])
cv2.imwrite(
os.path.join(save_path,
name[0] + '_' + str(idx) + '_maskv.png'),
fm_table[idx][1])
dir_path = r'D:\a上财学习材料\20毕业论文\实验\data'
filename_path = 'cTDaR19_imagename.txt'
write_filename(os.path.join(dir_path, filename_path), image_name)
# split train and test
split_train_val(dir_path, filename_path)
def train_net(net,
epochs=5,
batch_size=1,
lr=0.003,
val_percent=0.20,
loss_lambda=5,
save_cp=True,
gpu=False,
img_scale=0.5,
expositions_num=15,
logg_freq=15,
tb=False,
w_decay=0.0005,
use_notifications=False,
polyaxon=False,
outputs_path='checkpoints'):
# === Localize training data ===================================================
if polyaxon:
data_paths = get_data_paths()
dir_checkpoints = get_outputs_path()
dataSets_dir = os.path.join(data_paths['data1'], 'eprado',
'USLDR-DataSet')
#dataSets_dir = os.path.join(data_paths['data1'] , 'eprado', 'LDR_DataSet')
else:
dataSets_dir = os.path.join(wk_dir, "LDR_DataSet")
dir_checkpoints = os.path.join(wk_dir, outputs_path)
print('Dataset_dir', dataSets_dir)
print('Outputs_path', dir_checkpoints)
experiment_id = datetime.datetime.now().strftime('%d%m_%H%M_')
experiment_name = 'ExpandnetL_psn_{}_bs{}_lr{}_exps{}'.format(
experiment_id, batch_size, lr, expositions_num)
dir_img = os.path.join(dataSets_dir, 'Org_images/')
dir_compressions = os.path.join(dataSets_dir, 'c_images/')
dir_mask = os.path.join(dataSets_dir, 'c_images/')
#if tb:
#dummy_input = torch.rand(1, 3, 128, 128)
#writer.add_graph(net, (dummy_input,))
#writer.close()
# === Load Training/Validation data =====================================================
ids = get_ids(dir_compressions)
# Split into train test
idsset = list(ids)
kf = KFold(n_splits=5, shuffle=False)
#print('Train splits: ',kf.get_n_splits(dataset))
best_psnr_m = 0
best_psnr_hvs = 0
#for train_index, test_index in kf.split(idsset):
iddataset = split_train_val(idsset, expositions_num, val_percent)
#test_set = []
#for im_id in test_index:
# for e in range(expositions_num):
# test_set.append(idsset[im_id])
N_train = len(iddataset['train'])
N_val = len(iddataset['val'])
N_test = 0 #len(test_set)
#=====CHOOSE Loss Criterion=============================================================
#criterion = nn.MSELoss(reduction='mean')
criterion = ExpandNetLoss(loss_lambda=loss_lambda)
optimizer = optim.Adagrad(net.parameters(),
lr=lr,
lr_decay=0.000001,
weight_decay=w_decay)
#optimizer = optim.SGD(net.parameters(),
# lr=lr,
# momentum=0.9,
# weight_decay=0.0005)
since = time.time()
print('''
Training SETUP:
Epochs: {0:}
Batch size: {1:}
Optimizer: Adagrad
Learning rate: {2:}
Weight decay: {3:}
Training size: {4:}
Validation size: {5:}
Test size: {6:}
Checkpoints: {7:}
CUDA: {8:}
'''.format(epochs, batch_size, lr, w_decay, N_train, N_val, N_test,
str(save_cp), str(gpu)))
train_dataset = HdrDataset(iddataset['train'], dir_compressions, dir_mask,
expositions_num)
val_dataset = HdrDataset(iddataset['val'], dir_compressions, dir_mask,
expositions_num)
#test_dataset = HdrDataset(test_set, dir_compressions, dir_mask,expositions_num)
train_data_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
drop_last=False)
val_data_loader = DataLoader(val_dataset,
batch_size=batch_size,
shuffle=False,
drop_last=False)
#test_data_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=True)
best_hvsm = 0.0
global_psnr_m = []
global_psnr_hvs = []
for epoch in range(epochs):
print('\n')
print('{}{}{}'.format('+', '=' * 78, '+'))
print('| Starting epoch {}/{}. {}'.format(epoch + 1, epochs,
(' ' * 57) + '|'))
print('{}{}{}'.format('|', '-' * 78, '|'))
begin_of_epoch = time.time()
tot_steps = math.trunc(N_train / batch_size)
net.train()
train_loss = 0
losses = []
val_loss = 0
step = 0
train_sample = []
train_acc = 0
val_hvsm = 0
val_hvs = 0
model_pnsr_m = 0
for i, b in enumerate(train_data_loader):
step += 1
imgs, true_masks, imgs_ids = b['input'], b['target'], b['id']
#print(i, b['input'].size(), b['target'].size())
#input: [15, 3, 224, 224]), target: [15, 3, 224, 224]
#print('>>>>>>> Input max: ' , torch.max(imgs[0]))
#print('>>>>>>> mask max : ', torch.max(true_masks[0]))
if gpu:
imgs = imgs.cuda()
true_masks = true_masks.cuda()
else:
print(' GPU not available')
# Predicted mask images
optimizer.zero_grad()
prediction = net(imgs) #prediction shape: [B, 3, 224, 224]
#cost, cost_input_output = Hdr_loss(imgs, true_masks, prediction, sep_loss=False, gpu=gpu, tb=tb)
cost = criterion(prediction, true_masks)
#loss is torch tensor
losses.append(cost.item())
train_loss = np.mean(losses)
cost.backward()
optimizer.step()
if step == 1 or step % logg_freq == 0:
#print('| Step: {0:}, cost:{1:}, Train Loss:{2:.9f}, Train Acc:{3:.9f}'.format(step,cost, train_loss,train_acc/step))
print('| Step: {0:}, cost:{1:}, Train Loss:{2:.9f}'.format(
step, cost, train_loss))
#Last Step of this Epoch
if step == math.trunc(tot_steps):
num_in_batch = random.randrange(imgs.size(0))
train_sample_name = imgs_ids[num_in_batch]
train_sample = [
imgs[num_in_batch], true_masks[num_in_batch],
prediction[num_in_batch]
]
t_exp_name = 'Train_' + experiment_name
saveTocheckpoint(dir_checkpoints, t_exp_name,
train_sample_name, epoch, train_sample[0],
train_sample[1], train_sample[2])
if tb:
print(
'| saving train step {0:} sample : input,target & pred'
.format(step))
grid = torchvision.utils.make_grid(train_sample, nrow=3)
writer.add_image('train_sample', grid, 0)
#if epoch == 1 or epoch % 15 == 0 or epoch == epochs:
val_loss, val_hvsm, val_hvs = eval_hdr_net(net,
dir_checkpoints,
experiment_name,
val_data_loader,
criterion,
epoch,
gpu,
batch_size,
expositions_num=15,
tb=tb)
if tb:
writer.add_scalar('training_loss: ', train_loss, epoch)
writer.add_scalar('validation_loss', val_loss, epoch)
writer.add_scalar('val_hvsm', val_hvsm, epoch)
writer.add_scalar('val_hvs', val_hvs, epoch)
writer.add_scalars('losses', {
'training_loss': train_loss,
'val_loss': val_loss
}, epoch)
if polyaxon:
experiment.log_metrics(step=epoch,
training_loss=train_loss,
validation_loss=val_loss,
val_hvsm=val_hvsm,
val_hvs=val_hvs)
print('{}{}{}'.format('+', '=' * 78, '+'))
print('| {0:} Epoch {1:} finished ! {2:}|'.format(
' ' * 28, (epoch + 1), ' ' * 29))
print('{}{}{}'.format('+', '-' * 78, '+'))
print('| Summary: Train Loss: {0:0.07}, Val Loss:{1:}'.format(
train_loss, val_loss))
print('| Avrg psnr-hvs_m :{0:0.04},Avrg psnr-hvs :{1:0.04}'.
format(val_hvsm, val_hvs))
time_epoch = time.time() - begin_of_epoch
print('| Epoch ETC: {:.0f}m {:.0f}s'.format(time_epoch // 60,
time_epoch % 60))
print('{}{}{}'.format('+', '=' * 78, '+'))
if save_cp and (val_hvsm > best_hvsm):
best_hvsm = val_hvsm
model_path = os.path.join(dir_checkpoints, 'BestCP.pth')
torch.save(net.state_dict(), model_path)
print('Checkpoint saved !')
global_psnr_hvs.append(val_hvs)
global_psnr_m.append(val_hvsm)
'''
test_psnr_m, test_psnr_hvs = test_hdr_net(model_path,dir_checkpoints,
experiment_name,
test_data_loader,
criterion,gpu,tb)
if save_cp and (test_psnr_m > best_psnr_m):
best_psnr_m = test_psnr_m
best_model_path = os.path.join(dir_checkpoints, 'Best_CP.pth')
torch.save(net.state_dict(),best_model_path)
print('Best model saved !')
'''
print('>' * 80)
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Final Average psnr-hvs_m: {:.0f}, psnr-hvs: {:.0f}'.format(
np.mean(global_psnr_m), np.mean(global_psnr_hvs)))
if tb:
writer.close()
if use_notifications:
end_msg = "train.py finished at: {}(".format(
str(datetime.datetime.now()))
push = pb.push_note("usHDR: Finish", end_msg)
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,
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 train_primary(model_name, df_name, max_len, train_ratio, text_feature,
embeddings_version, embeddings_path, config_primary,
models_path, models_df):
device = utils.find_device()
loss_function, labels_num, dropout, epochs_num, threshold, \
hid_dim_lstm, lin_output_dim, lr, batch_size, momentum = \
utils.read_config_primary(config_primary, "train_primary")
train_df, val_df = utils.split_train_val(df_name, train_ratio)
train_dataloader, val_dataloader = utils.create_dataloaders_train(
train_df, val_df, text_feature, embeddings_version, max_len,
batch_size)
model = get_model(model_name, embeddings_path, hid_dim_lstm, loss_function,
labels_num, dropout, lin_output_dim)
model.to(device)
optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)
train_len = len(train_dataloader)
val_len = len(val_dataloader)
train_accuracy_list, train_loss_list, accuracy_list_val, loss_list_val = [], [], [], []
best_val_acc = 0.0
total_steps = len(train_dataloader) * epochs_num
scheduler = get_linear_schedule_with_warmup(optimizer,
num_warmup_steps=0,
num_training_steps=total_steps)
for epoch in range(epochs_num):
utils.print_epochs_progress(epoch, epochs_num)
start_train = time.clock()
acc_num_train = 0.0
loss_scalar = 0.0
model.train()
optimizer.zero_grad()
t0 = time.time()
predictions_dict_train = dict()
for batch_idx, batch in enumerate(train_dataloader):
if batch_idx % 10 == 0 and not batch_idx == 0:
utils.print_batches_progress(t0, batch_idx, train_dataloader)
input_ids = batch[0].to(device, dtype=torch.long)
masks = batch[1].to(device, dtype=torch.long)
labels = batch[2].to(device, dtype=torch.long)
hadm_id = batch[3]
model.zero_grad()
loss, predictions, probabilities = model(input_ids, masks, labels)
loss_scalar += loss.item()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
acc_num_train += utils.add_accuracy(predictions, labels)
optimizer.step()
scheduler.step()
predictions_dict_train = utils.update_predictions_dict(
predictions_dict_train, hadm_id, predictions, probabilities,
labels)
loss_scalar /= train_len
end_train = time.clock()
total_time = end_train - start_train
utils.print_train_epoch(epoch, acc_num_train, train_len, loss_scalar,
total_time)
train_loss_list.append(round(float(loss_scalar), 3))
train_accuracy_list.append(round(float(acc_num_train / train_len), 3))
utils.print_train_epoch_end(t0)
val_acc, val_loss, predictions_dict_val = evaluate_val(
model, val_dataloader, device, val_len)
accuracy_list_val.append(val_acc)
loss_list_val.append(val_loss)
if val_acc > best_val_acc:
torch.save(model.state_dict(), models_path + 'primary.pkl')
best_val_acc = val_acc
utils.save_predictions_to_df(predictions_dict_train, models_path,
'train', 'primary')
utils.save_predictions_to_df(predictions_dict_val, models_path,
'validation', 'primary')
utils.save_model(models_df, 'primary', hid_dim_lstm, labels_num,
loss_function, dropout, lin_output_dim, lr, epochs_num,
batch_size, momentum, best_val_acc)
utils.print_summary(models_path, 'primary', accuracy_list_val)
test_acc, predictions_dict_test = evaluate_test(model, device,
embeddings_version,
max_len, text_feature,
batch_size)
utils.save_predictions_to_df(predictions_dict_test, models_path, 'test',
'primary')
utils.print_test_results(models_path, 'primary', test_acc)
return models_df
if __name__ == "__main__":
# Data Loader
train_datagen = ImageDataGenerator(rotation_range=45,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.25,
horizontal_flip=True,
fill_mode='nearest')
test_datagen = ImageDataGenerator()
# 切分数据
(train_idx, valid_idx, ytr, yv, labels,
selected_breed_list) = split_train_val(labels, NUM_CLASSES, seed=SEED)
INPUT_SIZE = 299 # width/height of image in pixels (as expected by Xception model)
x_train = np.zeros((len(labels), INPUT_SIZE, INPUT_SIZE, 3),
dtype='float32')
for i, img_id in tqdm(enumerate(labels['id'])):
img = read_img(img_id, data_dir, 'train', (INPUT_SIZE, INPUT_SIZE))
x = xception.preprocess_input(np.expand_dims(img.copy(), axis=0))
x_train[i] = x
print('\nTotal Images shape: {}'.format(x_train.shape))
Xtr = x_train[train_idx]
Xv = x_train[valid_idx]
print('Train (images, H, W, C):', Xtr.shape, '\nVal (images, H, W, C):',
def train_net(net,
epochs=5,
batch_size=1,
lr=0.1,
val_percent=0.05,
save_cp=True,
gpu=False,
img_scale=1):
dir_img = 'data/train/'
dir_mask = 'data/train_masks/'
dir_checkpoint = 'checkpoints/'
ids = get_ids(dir_img)
iddataset = split_train_val(ids, val_percent)
traindata = BrainDataset(dir_img, dir_mask, iddataset['train'])
valdata = BrainDataset(dir_img, dir_mask, iddataset['val'])
trainloader = DataLoader(traindata,
batch_size=args.batchsize,
shuffle=True,
num_workers=4)
valloader = DataLoader(valdata, batch_size=1, shuffle=False, num_workers=2)
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()
epoch_tb = 0 # for tensorboard plotting
for epoch in range(epochs):
print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
if epoch % 4 == 0:
for param_group in optimizer.param_groups:
param_group['lr'] /= 2
print('learning rate == ' + str(param_group['lr']))
epoch_loss = 0
for i, data in enumerate(trainloader):
imgs = data[0]
true_masks = data[1]
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)
true_masks_flat = true_masks_flat.float()
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()))
writer.add_scalar('loss - iter', loss.item(), epoch_tb)
epoch_tb += 1
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch finished ! Loss: {}'.format(epoch_loss / (i + 1)))
writer.add_scalar('epoch_loss - epoch', epoch_loss / (i + 1), epoch)
if epoch % 7 == 0:
val_dice = eval_net(net, valloader, gpu, threshold=0.3)
print('Validation Dice Coeff 0.3: {}'.format(val_dice))
val_dice = eval_net(net, valloader, gpu, threshold=0.4)
print('Validation Dice Coeff 0.4: {}'.format(val_dice))
val_dice = eval_net(net, valloader, gpu, threshold=0.5)
print('Validation Dice Coeff 0.5: {}'.format(val_dice))
writer.add_scalar('val_dice_coeff - epoch', val_dice, epoch)
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.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_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(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,
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=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 = '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))
help='whether to save checkpoint')
(options, args) = parser.parse_args()
return options
if __name__ == '__main__':
args = get_args()
# dir_img = '/home/zzh/数据/mid project/raw_data'
# dir_mask = '/home/zzh/数据/mid project/groundtruth'
dir_img = '/home/zhuzhu/Desktop/mid project/raw_data'
dir_mask = '/home/zhuzhu/Desktop/mid project/groundtruth'
ids = get_ids(dir_img) # 1,2,3,...的生成器
iddataset = split_train_val(
ids, args.val_percent) # {'train':[23,98,59,...],'val':[12,37,48,...]}
net = UNet(n_channels=1, n_classes=args.num_classes)
optimizer = torch.optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.99,
weight_decay=5e-3)
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
criterion = nn.BCEWithLogitsLoss()
if args.load:
print('load model from checkpoint')
net.load_state_dict(torch.load('checkpoint/unet.pth'))
if args.gpu:
net.to(args.device)
if __name__ == '__main__':
import numpy as np
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(
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))
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.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,
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))
"gpu": args.gpu,
"load": args.load,
"channel_txt": args.channel_txt,
"scale": args.scale,
"lr": args.lr,
"iters": args.iters,
"epochs": args.epochs
},
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:
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,
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 get_dataloaders(meter_type, batch_size=20, trial=False):
train_records, test_records = get_train_test_records(meter_type)
with open(log_file, "a") as f:
f.write("{}: train = {}, test = {}\n".format("ABCDE"[meter_type],
len(train_records),
len(test_records)))
if meter_type == 2:
batch_size //= 2
pass
train_records, val_records = utils.split_train_val(train_records)
if trial:
n = 100
trainset = PointerMeters(meter_type,
train_records[:n],
transform=torchvision.transforms.Compose([
RandomCrop(meter_type),
Resize(meter_type),
ToTensor()
]))
valset = PointerMeters(meter_type,
val_records[:n],
transform=torchvision.transforms.Compose([
RandomCrop(meter_type),
Resize(meter_type),
ToTensor()
]))
else:
trainset = PointerMeters(meter_type,
train_records,
transform=torchvision.transforms.Compose([
RandomCrop(meter_type),
Resize(meter_type),
ToTensor()
]))
valset = PointerMeters(meter_type,
val_records,
transform=torchvision.transforms.Compose([
RandomCrop(meter_type),
Resize(meter_type),
ToTensor()
]))
train_loader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True)
val_loader = torch.utils.data.DataLoader(valset,
batch_size=batch_size,
shuffle=False)
with open(log_file, "a") as f:
f.write("train: {}, val: {}\n".format(len(train_loader.dataset),
len(val_loader.dataset)))
show_samples(train_loader, "train")
if trial:
np.random.shuffle(test_records)
testset = PointerMeters(meter_type,
test_records[:n],
transform=torchvision.transforms.Compose([
CentralCrop(meter_type),
Resize(meter_type),
ToTensor()
]))
else:
testset = PointerMeters(meter_type,
test_records,
transform=torchvision.transforms.Compose([
CentralCrop(meter_type),
Resize(meter_type),
ToTensor()
]))
test_loader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=True)
with open(log_file, "a") as f:
f.write("test: {}\n".format(len(test_loader.dataset)))
show_samples(test_loader, "test")
return train_loader, val_loader, test_loader
