def get_fcn_model(num_classes, use_gpu):
vgg_model = VGGNet(requires_grad=True, remove_fc=True, batch_norm=True)
fcn_model = FCN8sScaledBN(pretrained_net=vgg_model, n_class=num_classes)
if use_gpu:
ts = time.time()
vgg_model = vgg_model.cuda()
fcn_model = fcn_model.cuda()
num_gpu = list(range(torch.cuda.device_count()))
fcn_model = nn.DataParallel(fcn_model, device_ids=num_gpu)
print("Finish cuda loading, time elapsed {}".format(time.time() - ts))
return fcn_model
#
# val_data = Cityscapes(root, split='val', mode='fine',
# target_type='semantic', transform=trans, target_transform=trans)
# val_loader = DataLoader(val_data, batch_size=batch_size,
# shuffle=True, num_workers=0)
# img, show = train_data[0]
# print(show)
#############################################
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
fcn_model = FCNs(pretrained_net=vgg_model, n_class=n_class)
if use_gpu:
ts = time.time()
vgg_model = vgg_model.cuda()
fcn_model = fcn_model.cuda()
fcn_model = nn.DataParallel(fcn_model, device_ids=num_gpu)
print("Finish cuda loading, time elapsed {}".format(time.time() - ts))
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.RMSprop(fcn_model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=w_decay)
scheduler = lr_scheduler.StepLR(
optimizer, step_size=step_size,
gamma=gamma) # decay LR by a factor of 0.5 every 30 epochs
# create dir for score
score_dir = os.path.join("scores", configs)