def main():
# loading data
dataset = datasets.PixelLinkIC15Dataset(opt.train_images_dir,
opt.train_labels_dir)
sampler = WeightedRandomSampler([1 / len(dataset)] * len(dataset),
opt.batch_size,
replacement=True)
dataloader = DataLoader(dataset,
batch_size=opt.batch_size,
sampler=sampler)
my_net = net.Net() # construct neural network
# choose gpu or cpu
if opt.gpu:
device = torch.device("cuda:0")
my_net = my_net.cuda()
if opt.multi_gpu:
my_net = nn.DataParallel(my_net)
else:
device = torch.device("cpu")
# train, optimize
my_net.apply(weight_init)
optimizer = optim.SGD(my_net.parameters(),
lr=opt.learning_rate,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
optimizer2 = optim.SGD(my_net.parameters(),
lr=opt.learning_rate2,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
iteration = 0
train(opt.epoch, iteration, dataloader, my_net, optimizer, optimizer2,
device)
def retrain():
dataset = datasets.PixelLinkIC15Dataset(config.train_images_dir,
config.train_labels_dir)
sampler = WeightedRandomSampler([1 / len(dataset)] * len(dataset),
config.batch_size,
replacement=True)
dataloader = DataLoader(dataset,
batch_size=config.batch_size,
sampler=sampler)
my_net = net.Net()
if config.gpu:
device = torch.device("cuda:0")
my_net = my_net.cuda()
if config.multi_gpu:
my_net = nn.DataParallel(my_net)
else:
device = torch.device("cpu")
my_net.load_state_dict(
torch.load(config.saving_model_dir +
'%d.mdl' % config.retrain_model_index))
optimizer = optim.SGD(my_net.parameters(), lr=config.retrain_learning_rate2, \
momentum=config.momentum, weight_decay=config.weight_decay)
optimizer2 = optim.SGD(my_net.parameters(), lr=config.retrain_learning_rate, \
momentum=config.momentum, weight_decay=config.weight_decay)
train(config.retrain_epoch, config.retrain_model_index, dataloader, my_net,
optimizer, optimizer2, device)
def test_on_train_dataset(vis_per_img=10):
dataset = datasets.PixelLinkIC15Dataset(config.train_images_dir,
config.train_labels_dir,
train=False)
# dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=False)
my_net = net.Net()
if config.gpu:
device = torch.device("cuda:0")
my_net = my_net.cuda()
if config.multi_gpu:
my_net = nn.DataParallel(my_net)
else:
device = torch.device("cpu")
my_net.load_state_dict(
torch.load(config.saving_model_dir +
'%d.mdl' % config.test_model_index))
true_pos, true_neg, false_pos, false_neg = [0] * 4
for i in range(len(dataset)):
sample = dataset[i]
image = sample['image'].to(device)
image = image.unsqueeze(0)
my_labels = cal_label_on_batch(my_net, image)[0]
# print("my labels num: %d" % len(my_labels))
res = comp_gt_and_output(my_labels, sample["label"], 0.5)
if i % vis_per_img == 0:
image = image.squeeze(0).cpu().numpy()
image = ImgFormat.ImgOrderFormat(image,
from_order="CHW",
to_order="HWC")
image = ImgTransform.UnzeroMeanImage(image, config.r_mean,
config.g_mean, config.b_mean)
image = ImgFormat.ImgColorFormat(image,
from_color="RGB",
to_color="BGR")
image = visualize_label(image, my_labels, color=(0, 255, 0))
image = visualize_label(image,
sample["label"]["coor"],
color=(255, 0, 0))
cv2.imwrite("test_output/img_%d.jpg" % i, image)
true_pos += res[0]
false_pos += res[1]
false_neg += res[2]
if (true_pos + false_pos) > 0:
precision = true_pos / (true_pos + false_pos)
else:
precision = 0
if (true_pos + false_neg) > 0:
recall = true_pos / (true_pos + false_neg)
else:
recall = 0
print("i: %d, TP: %d, FP: %d, FN: %d, precision: %f, recall: %f" %
(i, true_pos, false_pos, false_neg, precision, recall))
def main():
dataset = datasets.PixelLinkIC15Dataset(config.train_images_dir,
config.train_labels_dir)
sampler = WeightedRandomSampler([1 / len(dataset)] * len(dataset),
config.batch_size,
replacement=True)
dataloader = DataLoader(dataset,
batch_size=config.batch_size,
sampler=sampler)
# dataloader = DataLoader(dataset, config.batch_size, shuffle=True)
my_net = net.Net()
if config.gpu:
device = torch.device("cuda:0")
my_net = my_net.cuda()
if config.multi_gpu:
my_net = nn.DataParallel(my_net)
else:
device = torch.device("cpu")
# nn.init.xavier_uniform_(list(my_net.parameters()))
my_net.apply(weight_init)
optimizer = optim.SGD(my_net.parameters(),
lr=config.learning_rate,
momentum=config.momentum,
weight_decay=config.weight_decay)
# if args.change:
optimizer2 = optim.SGD(my_net.parameters(),
lr=config.learning_rate2,
momentum=config.momentum,
weight_decay=config.weight_decay)
# else:
# optimizer2 = optim.SGD(my_net.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay)
iteration = 0
train(config.epoch, iteration, dataloader, my_net, optimizer, optimizer2,
device)
def main(retrain=False, test_datasets={}, vis_per_img=10):
res_dir = os.path.join(out_dir, 'snapshots')
if not os.path.exists(res_dir):
os.makedirs(res_dir)
# shutil.copyfile(os.path.join('configs', '%s.py' % exp_name), os.path.join(out_dir, 'config.py'))
with open(os.path.join(out_dir, 'config.py'), 'w') as f:
params_names = config.__dir__()
for param_name in params_names:
if param_name.startswith('__'):
continue
param_value = getattr(config, param_name)
if isinstance(param_value, str):
f.write('%s = "%s"\n' % (param_name, param_value))
else:
f.write('%s = %s\n' % (param_name, param_value))
dataset = datasets.PixelLinkIC15Dataset(
config.train_images_dir,
config.train_labels_dir,
all_trains=config.all_trains,
version=config.version,
mean=config.mean,
use_rotate=config.use_rotate,
use_crop=config.use_crop,
image_size_train=config.image_size_train,
image_size_test=config.image_size_test)
# sampler = WeightedRandomSampler([1/len(dataset)]*len(dataset), config.batch_size, replacement=True)
# dataloader = DataLoader(dataset, batch_size=config.batch_size, sampler=sampler)
dataloader = DataLoader(dataset,
config.batch_size,
shuffle=True,
num_workers=6)
model = net.PixelLinkNet(
**config.net_params) #net.Net(config.version, config.dilation)
if config.gpu:
device = torch.device("cuda:0")
model = model.cuda()
if config.multi_gpu:
model = nn.DataParallel(model)
else:
device = torch.device("cpu")
loss = PixelLinkLoss(config.pixel_weight, config.link_weight,
config.neg_pos_ratio)
optimizer = optim.SGD(model.parameters(),
lr=config.learning_rate1,
momentum=config.momentum,
weight_decay=config.weight_decay)
epoch_milestone = math.ceil(config.step2_start / len(dataloader))
print('LR schedule')
print('[%05d - %05d] : %E' % (0, epoch_milestone, config.learning_rate1))
print('[%05d - %05d] : %E' %
(epoch_milestone, config.epoch, config.learning_rate2))
scheduler = optim.lr_scheduler.MultiStepLR(
optimizer, [epoch_milestone],
config.learning_rate2 / config.learning_rate1)
global trainer, callbacks_cont
tqdm = callbacks.TQDM()
log_path = os.path.join(out_dir, 'log_train.csv')
index = 0
while os.path.exists(log_path):
index += 1
log_path = os.path.join(out_dir, 'log_train_%02d.csv' % index)
logger = callbacks.CSVLogger(log_path)
trainer = ModuleTrainer(model)
trainer.compile(optimizer, loss, callbacks=[tqdm, logger])
callbacks_cont = callbacks.CallbackContainer(trainer._callbacks)
callbacks_cont.set_trainer(trainer)
if retrain:
# find latest snapshot
snapshots_dir = os.path.join(out_dir, 'snapshots')
model_files = glob.glob(snapshots_dir + '/epoch_*')
if model_files:
resume_path = sorted(model_files)[-1]
start_epoch = 1 + int(
os.path.basename(resume_path)[len('epoch_'):-4])
print('Loading snapshot from : %s' % resume_path)
checkpoint = torch.load(resume_path)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
# couldnt find snapshots
start_epoch = 0
else:
start_epoch = 0
train(config.epoch,
dataloader,
model,
loss,
optimizer,
scheduler,
device,
start_epoch=start_epoch,
test_datasets=test_datasets,
vis_per_img=vis_per_img)
model,
loss,
optimizer,
scheduler,
device,
start_epoch=start_epoch,
test_datasets=test_datasets,
vis_per_img=vis_per_img)
if __name__ == "__main__":
if args.mode in ['train', 'retrain']:
dataset_train = datasets.PixelLinkIC15Dataset(
config.train_images_dir,
config.train_labels_dir,
train=False,
all_trains=config.all_trains,
version=config.version,
mean=config.mean,
image_size_test=config.image_size_test)
vis_per_img_train = int(math.ceil(config.all_trains / 100.0))
dataset_test = datasets.PixelLinkIC15Dataset(
config.test_images_dir,
config.test_labels_dir,
train=False,
all_trains=config.all_tests,
version=config.version,
mean=config.mean,
image_size_test=config.image_size_test)
vis_per_img_test = int(math.ceil(config.all_tests / 100.0))