def main():
net = DenseASPP(model_cfg=DenseASPP121.Model_CFG).cuda()
# densenet121 = models.densenet121(pretrained=True)
if len(args.checkpoint_path) == 0:
curr_epoch = 1
# Initializing 'best_record'
args.best_record = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
# load the pretrained model
print('training resumes from ' + args.checkpoint_path)
# lambda argument: manipulate(argument)
# pretrained_weight = torch.load(args.checkpoint_path, map_location=lambda storage, loc: storage)
pretrained_weight = torch.load(args.checkpoint_path)
new_state_dict = OrderedDict()
model_dict = net.state_dict()
for key, value in pretrained_weight.items():
name = key
# print(name)
new_state_dict[name] = value
new_state_dict.pop('features.norm5.weight')
new_state_dict.pop('features.norm5.bias')
new_state_dict.pop('features.norm5.running_mean')
new_state_dict.pop('features.norm5.running_var')
new_state_dict.pop('classifier.weight')
new_state_dict.pop('classifier.bias')
# print(new_state_dict)
model_dict.update(new_state_dict)
net.load_state_dict(model_dict, strict=False)
# pretrained_dict = {key: value for key, value in pretrained_dict.items() if key in model_dict}
# model_dict.update(pretrained_dict)
# pretrained_dict = {key: value for key, value in pretrained_dict.items() if key != 'classifier.weight' or 'classifier.bias'}
# for key, value in pretrained_dict.items():
# if 'classifier.weight' in key:
# key = key.rstrip('classifier.weight')
# if 'classifier.bias' in key:
# key = key.rstrip('classifier.bias')
# if 'features.norm5.weight' in key:
# key = key.replace("features.norm5.weight", "")
# elif 'features.norm5.bias' in key:
# key = key.strip('features.norm5.bias')
# elif 'features.norm5.running_mean' in key:
# key = key.strip('features.norm5.running_mean')
# elif 'features.norm5.running_var' in key:
# key = key.strip('features.norm5.running_var')
# name = key
# print(name)
# new_pretrained_dict[name] = value
# model.load_state_dict(model_dict, strict=False)
# model.load_state_dict(new_pretrained_dict, strict=False)
curr_epoch = 1
args.best_record = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
# split_snapshot = args.checkpoint_path.split('_')
# curr_epoch = int(split_snapshot[1]) + 1
# args.best_record = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
# 'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
# 'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
cudnn.benchmark = True
# net.train()
# tells your model that you are training the model ,or sets the module in training mode.
# The classic workflow : call train() --> epoch of training on the training set --> call eval()
# --> evaluate your model on the validation set
# mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
mean_std = ([0.290101, 0.328081, 0.286964], [0.182954, 0.186566, 0.184475])
# ---------------------------------- [[ data - augmentation ]] ---------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------
# [[joint_transforms]]
# both raw image and gt are transformed by data-augmentation
train_joint_transform = joint_transforms.Compose([
joint_transforms.RandomSizedCrop(size=args.input_width),
# joint_transforms.RandomSized(size=args.input_width),
joint_transforms.RandomHorizontallyFlip()
])
val_joint_transform = joint_transforms.Compose([
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomSizedCrop(size=args.input_width)
])
# transform : To preprocess images
# Compose : if there are a lot of preprocessed images, compose plays a role as collector in a single space.
input_transform = standard_transforms.Compose([
standard_transforms.ColorJitter(hue=0.1),
# colorjitter.ColorJitter(brightness=0.1),
standard_transforms.ToTensor(),
# standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ColorJitter(hue=0.1),
standard_transforms.ToTensor(),
# standard_transforms.Normalize(*mean_std)
])
# target_transform = extended_transforms.MaskToTensor()
target_transform = extended_transforms.Compose([
extended_transforms.MaskToTensor(),
])
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
# [[ ToPILImage() ]]
# Convert a tensor or an ndarray to PIL Image.
standard_transforms.ToPILImage()
])
visualize = standard_transforms.ToTensor()
# -----------------------------------------------------------------------------------------------------------------
# -----------------------------------------------------------------------------------------------------------------
train_set = segmentation_dataloader.CityScapes(
'fine',
'train',
joint_transform=train_joint_transform,
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=args.train_batch_size,
num_workers=args.num_threads,
shuffle=True)
val_set = segmentation_dataloader.CityScapes(
'fine',
'val',
joint_transform=val_joint_transform,
transform=val_input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=args.val_batch_size,
num_workers=args.num_threads,
shuffle=False)
criterion = torch.nn.CrossEntropyLoss(
ignore_index=segmentation_dataloader.ignore_label)
optimizer = optim.Adam(net.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# optimizer = optim.Adam([
# { 'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
# 'lr': args.learning_rate}
# ], weight_decay=args.weight_decay)
# if len(args.checkpoint_path) > 0:
# optimizer.load_state_dict(torch.load(args.checkpoint_path))
# optimizer.param_groups[0]['lr'] = 2 * args.learning_rate
# optimizer.param_groups[1]['lr'] = args.learning_rate
check_mkdir(ckpt_path)
# check_mkdir(os.path.join(ckpt_path, exp_name))
check_mkdir(os.path.join(ckpt_path, 'Model', ImageNet, exp_name_ImageNet))
# open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(args) + '\n\n')
open(
os.path.join(ckpt_path, 'Model', ImageNet, exp_name_ImageNet,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(args) + '\n\n')
# lambda1 = lambda epoch : (1 - epoch // args['epoch_num']) ** 0.9
# [[learning-rate decay]]
# factor = (1 - curr_epoch/ args['epoch_num'])**0.9
# scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor= factor, patience=args['lr_patience'],
# min_lr=0)
for epoch in range(curr_epoch, args.num_epochs + 1):
# factor = (1 - epoch / args['epoch_num']) ** 0.9
# scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=factor, patience=args['lr_patience'],
# min_lr=0)
# [[ training ]]
train(train_loader, net, criterion, optimizer, epoch, args, train_set)
# train(train_loader, net, optimizer, epoch, args, train_set)
# [[ validation ]]
validate(val_loader, net, criterion, optimizer, epoch, args,
restore_transform, visualize)
# validate(val_loader, net, optimizer, epoch, args, restore_transform, visualize)
# scheduler.step(val_loss)
print('Training Done!!')
def main():
net = DenseASPP_boundary_depthwise(model_cfg=DenseASPP121.Model_CFG).cuda()
# net = DenseASPP(model_cfg=DenseASPP121.Model_CFG).cuda()
# densenet121 = models.densenet121(pretrained=True)
if len(args.checkpoint_path) == 0:
curr_epoch = 1
# Initializing 'best_record'
args.best_record = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
# load the pretrained model
print('training resumes from ' + args.checkpoint_path)
# lambda ==> argument: manipulate(argument)
pretrained_weight = torch.load(
args.checkpoint_path, map_location=lambda storage, loc: storage)
""" map_location = lambda storage, loc: storage--> Load all tensors onto the CPU, using a function"""
new_state_dict = OrderedDict()
model_dict = net.state_dict()
for key, value in pretrained_weight.items():
name = key
new_state_dict[name] = value
if name.find('norm') >= 9:
# print('norm contained from pretrained_weight : ', name)
value.requires_grad = False
# if name.find('conv0') >= 9:
# print('norm contained from pretrained_weight : ', name)
# value.requires_grad = False
new_state_dict.pop('features.conv0.weight')
new_state_dict.pop('features.norm5.weight')
new_state_dict.pop('features.norm5.bias')
new_state_dict.pop('features.norm5.running_mean')
new_state_dict.pop('features.norm5.running_var')
new_state_dict.pop('classifier.weight')
new_state_dict.pop('classifier.bias')
model_dict.update(new_state_dict)
net.load_state_dict(model_dict)
# pretrained_dict = {key: value for key, value in pretrained_dict.items() if key in model_dict}
# model_dict.update(pretrained_dict)
# pretrained_dict = {key: value for key, value in pretrained_dict.items() if key != 'classifier.weight' or 'classifier.bias'}
# model.load_state_dict(model_dict, strict=False)
# model.load_state_dict(new_pretrained_dict, strict=False)
curr_epoch = 1
args.best_record = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
# ---------------------------------- [[ data - augmentation ]] ---------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------
# [[joint_transforms]]
# both raw image and gt are transformed by data-augmentation
train_joint_transform = joint_transforms.Compose([
# joint_transforms.ImageScaling(size=[0.5, 2.0]),
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomSizedCrop(size=args.input_width),
])
# transform : To preprocess images
# Compose : if there are a lot of preprocessed images, compose plays a role as collector in a single space.
input_transform = standard_transforms.Compose([
# Colorjitter.ColorJitter(brightness=[-10, 10]),
standard_transforms.ColorJitter(hue=0.1),
standard_transforms.ToTensor(),
# standard_transforms.Normalize(*my_mean_std)
])
target_transform = extended_transforms.MaskToTensor()
train_set = segmentation_dataloader.CityScapes(
'fine',
'train',
joint_transform=train_joint_transform,
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=args.train_batch_size,
num_workers=args.num_threads,
shuffle=True)
# optimizer = optim.Adam(net.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
criterion = torch.nn.CrossEntropyLoss(
ignore_index=segmentation_dataloader.ignore_label).cuda()
num_training_samples = len(train_set)
steps_per_epoch = np.ceil(num_training_samples /
args.train_batch_size).astype(np.int32)
num_total_steps = args.num_epochs * steps_per_epoch
print("total number of samples: {}".format(num_training_samples))
print("total number of steps : {}".format(num_total_steps))
# COUNT_PARAMS
total_num_paramters = 0
for param in net.parameters():
total_num_paramters += np.array(list(param.size())).prod()
print("number of trainable parameters: {}".format(total_num_paramters))
for epoch in range(curr_epoch, args.num_epochs + 1):
lr_ = poly_lr_scheduler(init_lr=args.learning_rate, epoch=epoch - 1)
optimizer = optim.Adam(net.parameters(),
lr=lr_,
weight_decay=args.weight_decay)
train(train_loader, net, criterion, optimizer, epoch, args,
total_num_paramters)
torch.save(
net.state_dict(),
os.path.join(ckpt_path, 'Model', ImageNet, exp_name_ImageNet,
'model-{}'.format(total_num_paramters) + '.pkl'))
print('Training Done!!')
def main():
net = DenseASPP(model_cfg=DenseASPP121.Model_CFG).cuda()
# densenet121 = models.densenet121(pretrained=True)
if len(args.checkpoint_path) == 0:
curr_epoch = 1
# Initializing 'best_record'
args.best_record = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
# load the pretrained model
print('training resumes from ' + args.checkpoint_path)
# lambda ==> argument: manipulate(argument)
pretrained_weight = torch.load(
args.checkpoint_path, map_location=lambda storage, loc: storage)
""" map_location = lambda storage, loc: storage--> Load all tensors onto the CPU, using a function"""
new_state_dict = OrderedDict()
model_dict = net.state_dict()
for key, value in pretrained_weight.items():
name = key
new_state_dict[name] = value
if name.find('norm') >= 9:
print('norm contained from pretrained_weight : ', name)
value.requires_grad = False
# if name.find('conv0') >= 9:
# print('norm contained from pretrained_weight : ', name)
# value.requires_grad = False
new_state_dict.pop('features.norm5.weight')
new_state_dict.pop('features.norm5.bias')
new_state_dict.pop('features.norm5.running_mean')
new_state_dict.pop('features.norm5.running_var')
new_state_dict.pop('classifier.weight')
new_state_dict.pop('classifier.bias')
model_dict.update(new_state_dict)
net.load_state_dict(model_dict)
# pretrained_dict = {key: value for key, value in pretrained_dict.items() if key in model_dict}
# model_dict.update(pretrained_dict)
# pretrained_dict = {key: value for key, value in pretrained_dict.items() if key != 'classifier.weight' or 'classifier.bias'}
# model.load_state_dict(model_dict, strict=False)
# model.load_state_dict(new_pretrained_dict, strict=False)
curr_epoch = 1
args.best_record = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
# cudnn.benchmark = True
# net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
my_mean_std = ([0.688824, 0.270588,
0.305882], [0.041225, 0.032375, 0.026271])
# ---------------------------------- [[ data - augmentation ]] ---------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------
# [[joint_transforms]]
# both raw image and gt are transformed by data-augmentation
train_joint_transform = joint_transforms.Compose([
# joint_transforms.ImageScaling(size=[0.5, 2.0]),
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomSizedCrop(size=args.input_width),
])
val_joint_transform = joint_transforms.Compose([
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomSizedCrop(size=args.input_width),
])
"""random_gamma = tf.random_uniform([], 0.8, 1.2)
left_image_aug = left_image ** random_gamma
right_image_aug = right_image ** random_gamma
# randomly shift brightness
random_brightness = tf.random_uniform([], 0.5, 2.0)
left_image_aug = left_image_aug * random_brightness
right_image_aug = right_image_aug * random_brightness"""
# transform : To preprocess images
# Compose : if there are a lot of preprocessed images, compose plays a role as collector in a single space.
input_transform = standard_transforms.Compose([
# Colorjitter.ColorJitter(brightness=[-10, 10]),
standard_transforms.ColorJitter(hue=0.1),
standard_transforms.ToTensor(),
# standard_transforms.Normalize(*my_mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ColorJitter(hue=0.1),
standard_transforms.ToTensor(),
# standard_transforms.Normalize(*my_mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
# """[[ ToPILImage() ]]"""
# Convert a tensor or an ndarray to PIL Image.
standard_transforms.ToPILImage()
])
visualize = standard_transforms.ToTensor()
# -----------------------------------------------------------------------------------------------------------------
# -----------------------------------------------------------------------------------------------------------------
train_set = segmentation_dataloader.CityScapes(
'fine',
'train',
joint_transform=train_joint_transform,
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=args.train_batch_size,
num_workers=args.num_threads,
shuffle=True)
val_set = segmentation_dataloader.CityScapes(
'fine',
'val',
joint_transform=val_joint_transform,
transform=val_input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=args.val_batch_size,
num_workers=args.num_threads,
shuffle=False)
# optimizer = optim.Adam(net.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
criterion = torch.nn.CrossEntropyLoss(
ignore_index=segmentation_dataloader.ignore_label).cuda()
for epoch in range(curr_epoch, args.num_epochs + 1):
# net.train()
lr_ = poly_lr_scheduler(init_lr=args.learning_rate, epoch=epoch - 1)
optimizer = optim.Adam(net.parameters(),
lr=lr_,
weight_decay=args.weight_decay)
train(train_loader, net, criterion, optimizer, epoch, args, train_set)
net.eval()
validate(val_loader, net, criterion, optimizer, epoch, args,
restore_transform, visualize)
print('Training Done!!')