def main(args):
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.method))
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
# conduct seg network
seg_model = get_model(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = seg_model.state_dict().copy()
if args.init:
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0] == 'fc':
new_params['encoder.' + '.'.join(i_parts[:])] = saved_state_dict[i]
seg_model.load_state_dict(new_params)
print('loading params w/o fc')
else:
seg_model.load_state_dict(saved_state_dict)
print('loading params all')
model = DataParallelModel(seg_model)
model.float()
model.cuda()
# define dataloader
train_loader = data.DataLoader(DataGenerator(root=args.root, list_path=args.lst,
crop_size=args.crop_size, training=True),
batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
val_loader = data.DataLoader(DataGenerator(root=args.val_root, list_path=args.val_lst,
crop_size=args.crop_size, training=False),
batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True)
# define criterion & optimizer
criterion = ABRLovaszLoss(ignore_index=args.ignore_label, only_present=True, cls_p= args.num_classes, cls_h= args.hbody_cls, cls_f= args.fbody_cls)
criterion = DataParallelCriterion(criterion).cuda()
optimizer = optim.SGD(
[{'params': filter(lambda p: p.requires_grad, seg_model.parameters()), 'lr': args.learning_rate}],
lr=args.learning_rate, momentum=0.9, weight_decay=5e-4)
# key points
best_val_mIoU = 0
best_val_pixAcc = 0
start = time.time()
for epoch in range(0, args.epochs):
print('\n{} | {}'.format(epoch, args.epochs - 1))
# training
_ = train(model, train_loader, epoch, criterion, optimizer, writer)
# validation
if epoch %10 ==0 or epoch > args.epochs*0.8:
val_pixacc, val_miou = validation(model, val_loader, epoch, writer)
# save model
if val_pixacc > best_val_pixAcc:
best_val_pixAcc = val_pixacc
if val_miou > best_val_mIoU:
best_val_mIoU = val_miou
model_dir = os.path.join(args.snapshot_dir, args.method + '_miou.pth')
torch.save(seg_model.state_dict(), model_dir)
print('Model saved to %s' % model_dir)
os.rename(model_dir, os.path.join(args.snapshot_dir, args.method + '_miou'+str(best_val_mIoU)+'.pth'))
print('Complete using', time.time() - start, 'seconds')
print('Best pixAcc: {} | Best mIoU: {}'.format(best_val_pixAcc, best_val_mIoU))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
config = yaml.load(open('configs/config.yaml'), Loader=yaml.FullLoader)
train_set, val_set = tilegenerator(image_paths=config['train']['folder'],
num_images=config['train']['num_images'],
train_val_ratio=0.9)
image_datasets = {'train': train_set, 'val': val_set}
batch_size = config['train']['batch_size']
dataset = {
'train':
DataLoader(DataGenerator(train_set),
batch_size=batch_size,
shuffle=True,
num_workers=config['train']['num_workers']),
'val':
DataLoader(DataGenerator(val_set),
batch_size=4,
shuffle=True,
num_workers=0)
}
dataset_sizes = {x: len(image_datasets[x]) for x in image_datasets.keys()}
logging.info(dataset_sizes)
model = ResNet50().cuda()