def __init__(self, args, config, cuda=None):
self.args = args
os.environ["CUDA_VISIBLE_DEVICES"] = self.args.gpu
self.config = config
self.cuda = cuda and torch.cuda.is_available()
self.device = torch.device('cuda' if self.cuda else 'cpu')
self.best_MIou = 0
self.current_epoch = 0
self.epoch_num = self.config.epoch_num
self.current_iter = 0
self.writer = SummaryWriter()
# path definition
self.val_list_filepath = os.path.join(
args.data_root_path, 'VOC2012/ImageSets/Segmentation/val.txt')
self.gt_filepath = os.path.join(args.data_root_path,
'VOC2012/SegmentationClass/')
self.pre_filepath = os.path.join(args.data_root_path,
'VOC2012/JPEGImages/')
# Metric definition
self.Eval = Eval(self.config.num_classes)
# loss definition
if args.loss_weight:
classes_weights_path = os.path.join(
self.config.classes_weight,
self.args.dataset + 'classes_weights_log.npy')
print(classes_weights_path)
if not os.path.isfile(classes_weights_path):
logger.info('calculating class weights...')
calculate_weigths_labels(self.config)
class_weights = np.load(classes_weights_path)
pprint.pprint(class_weights)
weight = torch.from_numpy(class_weights.astype(np.float32))
logger.info('loading class weights successfully!')
else:
weight = None
self.loss = nn.CrossEntropyLoss(weight=weight, ignore_index=255)
self.loss.to(self.device)
# model
self.model = DeepLab(output_stride=self.args.output_stride,
class_num=self.config.num_classes,
pretrained=self.args.imagenet_pretrained,
bn_momentum=self.args.bn_momentum,
freeze_bn=self.args.freeze_bn)
self.model = nn.DataParallel(self.model, device_ids=range(4))
patch_replication_callback(self.model)
self.model.to(self.device)
self.optimizer = torch.optim.SGD(
params=[
{
"params": self.get_params(self.model.module, key="1x"),
"lr": self.args.lr,
},
{
"params": self.get_params(self.model.module, key="10x"),
"lr": 10 * self.args.lr,
},
],
momentum=self.config.momentum,
# dampening=self.config.dampening,
weight_decay=self.config.weight_decay,
# nesterov=self.config.nesterov
)
# dataloader
self.dataloader = VOCDataLoader(self.args, self.config)
def __init__(self, args, cuda=None):
self.args = args
os.environ["CUDA_VISIBLE_DEVICES"] = self.args.gpu
self.cuda = cuda and torch.cuda.is_available()
self.device = torch.device('cuda' if self.cuda else 'cpu')
self.current_MIoU = 0
self.best_MIou = 0
self.current_epoch = 0
self.current_iter = 0
# set TensorboardX
self.writer = SummaryWriter(log_dir=self.args.run_name)
# Metric definition
self.Eval = Eval(self.args.num_classes)
# loss definition
if self.args.loss_weight_file is not None:
classes_weights_path = os.path.join(self.args.loss_weights_dir,
self.args.loss_weight_file)
print(classes_weights_path)
if not os.path.isfile(classes_weights_path):
logger.info('calculating class weights...')
calculate_weigths_labels(self.args)
class_weights = np.load(classes_weights_path)
pprint.pprint(class_weights)
weight = torch.from_numpy(class_weights.astype(np.float32))
logger.info('loading class weights successfully!')
else:
weight = None
self.loss = nn.CrossEntropyLoss(weight=weight, ignore_index=255)
self.loss.to(self.device)
# model
self.model = Unet_decoder(output_stride=self.args.output_stride,
class_num=self.args.num_classes,
pretrained=self.args.imagenet_pretrained
and self.args.pretrained_ckpt_file == None,
bn_momentum=self.args.bn_momentum,
freeze_bn=self.args.freeze_bn)
self.model = nn.DataParallel(self.model,
device_ids=range(
ceil(len(self.args.gpu) / 2)))
patch_replication_callback(self.model)
self.model.to(self.device)
self.optimizer = torch.optim.SGD(
params=[
{
"params": self.get_params(self.model.module, key="1x"),
"lr": self.args.lr,
},
{
"params": self.get_params(self.model.module, key="10x"),
"lr": 10 * self.args.lr,
},
],
momentum=self.args.momentum,
# dampening=self.args.dampening,
weight_decay=self.args.weight_decay,
# nesterov=self.args.nesterov
)
# dataloader
self.dataloader = VOCDataLoader(self.args)
self.epoch_num = ceil(self.args.iter_max /
self.dataloader.train_iterations)
