def __init__(self, mode):
# Define Saver
self.saver = Saver(opt, mode)
self.logger = self.saver.logger
# Visualize
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Dataset dataloader
self.train_dataset, self.train_loader = make_data_loader(opt)
self.nbatch_train = len(self.train_loader)
self.val_dataset, self.val_loader = make_data_loader(opt, mode="val")
self.nbatch_val = len(self.val_loader)
# Model
if opt.sync_bn is None and len(opt.gpu_id) > 1:
opt.sync_bn = True
else:
opt.sync_bn = False
# model = DeepLab(opt)
# model = CSRNet()
model = CRGNet(opt)
model_info(model, self.logger)
self.model = model.to(opt.device)
# Loss
if opt.use_balanced_weights:
classes_weights_file = osp.join(opt.root_dir, 'train_classes_weights.npy')
if os.path.isfile(classes_weights_file):
weight = np.load(classes_weights_file)
else:
weight = calculate_weigths_labels(
self.train_loader, opt.root_dir)
print(weight)
opt.loss['weight'] = weight
self.loss = build_loss(opt.loss)
# Define Evaluator
self.evaluator = Evaluator() # use region to eval: class_num is 2
# Resuming Checkpoint
self.best_pred = 0.0
self.start_epoch = 0
if opt.resume:
if os.path.isfile(opt.pre):
print("=> loading checkpoint '{}'".format(opt.pre))
checkpoint = torch.load(opt.pre)
self.start_epoch = checkpoint['epoch']
self.best_pred = checkpoint['best_pred']
self.model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(opt.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(opt.pre))
if len(opt.gpu_id) > 1:
print("Using multiple gpu")
self.model = torch.nn.DataParallel(self.model,
device_ids=opt.gpu_id)
# Define Optimizer
# train_params = [{'params': model.get_1x_lr_params(), 'lr': opt.lr},
# {'params': model.get_10x_lr_params(), 'lr': opt.lr * 10}]
# self.optimizer = torch.optim.SGD(train_params,
# momentum=opt.momentum,
# weight_decay=opt.decay)
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=opt.lr,
momentum=opt.momentum,
weight_decay=opt.decay)
# Define lr scheduler
# self.scheduler = LR_Scheduler(mode=opt.lr_scheduler,
# base_lr=opt.lr,
# num_epochs=opt.epochs,
# iters_per_epoch=self.nbatch_train,
# lr_step=140)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer,
milestones=[round(opt.epochs * x) for x in opt.steps],
gamma=opt.gamma)
# Time
self.loss_hist = collections.deque(maxlen=500)
self.timer = Timer(opt.epochs, self.nbatch_train, self.nbatch_val)
self.step_time = collections.deque(maxlen=opt.print_freq)
def __init__(self, mode):
# Define Saver
self.saver = Saver(opt, mode)
self.logger = self.saver.logger
# visualize
self.summary = TensorboardSummary(self.saver.experiment_dir, opt)
self.writer = self.summary.writer
# Define Dataloader
# train dataset
self.train_dataset, self.train_loader = make_data_loader(opt, train=True)
self.nbatch_train = len(self.train_loader)
self.num_classes = self.train_dataset.num_classes
# val dataset
self.val_dataset, self.val_loader = make_data_loader(opt, train=False)
self.nbatch_val = len(self.val_loader)
# Define Network
# initilize the network here.
self.model = Model(opt, self.num_classes)
self.model = self.model.to(opt.device)
# Detection post process(NMS...)
self.post_pro = PostProcess(**opt.nms)
# Define Optimizer
if opt.adam:
self.optimizer = optim.Adam(self.model.parameters(), lr=opt.lr)
else:
self.optimizer = optim.SGD(self.model.parameters(), lr=opt.lr, momentum=opt.momentum, weight_decay=opt.decay)
# Apex
if opt.use_apex:
self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level='O1')
# Resuming Checkpoint
self.best_pred = 0.0
self.start_epoch = 0
if opt.resume:
if os.path.isfile(opt.pre):
print("=> loading checkpoint '{}'".format(opt.pre))
checkpoint = torch.load(opt.pre)
self.start_epoch = checkpoint['epoch'] + 1
self.best_pred = checkpoint['best_pred']
self.model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(opt.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(opt.pre))
# Define lr scherduler
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
# self.optimizer, patience=3, verbose=True)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer,
milestones=[round(opt.epochs * x) for x in opt.steps],
gamma=opt.gamma)
self.scheduler.last_epoch = self.start_epoch - 1
# Using mul gpu
if len(opt.gpu_id) > 1:
self.logger.info("Using multiple gpu")
self.model = torch.nn.DataParallel(self.model,
device_ids=opt.gpu_id)
# metrics
if opt.eval_type == 'cocoeval':
self.eval = COCO_eval(self.val_dataset.coco)
else:
self.eval = VOC_eval(self.num_classes)
self.loss_hist = collections.deque(maxlen=500)
self.timer = Timer(opt.epochs, self.nbatch_train, self.nbatch_val)
self.step_time = collections.deque(maxlen=opt.print_freq)
parser.add_argument('--validate', action='store_true',
help='validate')
# checking point
parser.add_argument('--resume', type=str, default=None,
help='put the path to resuming file if needed')
parser.add_argument('--checkname', type=str, default='VesselNN_Unsupervised',
help='set the checkpoint name')
args = parser.parse_args()
# Define Saver
saver = Saver(args)
saver.save_experiment_config()
# Define Tensorboard Summary
summary = TensorboardSummary(saver.experiment_dir)
writer = summary.create_summary()
# Data
dataset = Directory_Image_Train(images_path=args.train_images_path,
labels_path=args.train_labels_path,
data_shape=(32, 128, 128),
lables_shape=(32, 128, 128),
range_norm=args.range_norm)
dataloader = DataLoader(dataset, batch_size=torch.cuda.device_count() * args.batch_size, shuffle=True, num_workers=2)
# Data - validation
dataset_val = Single_Image_Eval(image_path=args.val_image_path,
label_path=args.val_label_path,
data_shape=(32, 128, 128),
lables_shape=(32, 128, 128),
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(args.logdir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
dltrain = DLDataset('trainval', "./data/pascal_voc_seg/tfrecord/")
dlval = DLDataset('val', "./data/pascal_voc_seg/tfrecord/")
# dltrain = DLDataset('trainval', "./data/pascal_voc_seg/VOCdevkit/VOC2012/")
# dlval = DLDataset('val', "./data/pascal_voc_seg/VOCdevkit/VOC2012/")
self.train_loader = DataLoader(dltrain,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
self.val_loader = DataLoader(dlval,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
# Define network
model = Deeplab()
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
self.criterion = nn.CrossEntropyLoss(ignore_index=255).cuda()
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(21)
# Define lr scheduler
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer=optimizer)
# Using cuda
# if args.cuda:
# self.model = torch.nn.DataParallel(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0