num_workers=workers)
#val_dataset = CityscapesLoader(split = 'val')
val_loader = DataLoader(
val_dataset, batch_size=batch_size, shuffle=False, num_workers=workers)
num_class = {'SUN':37,'NYU':40,'CITY':19}[args.dataset]
ignore_label = {'SUN':255,'NYU':255,'CITY':250}[args.dataset]
loss = nn.CrossEntropyLoss(ignore_index=ignore_label)
patience = {'SUN':15,'NYU':60,'CITY':40}[args.dataset]
print('Train sample number: %d' % len(train_dataset))
print('Val sample number: %d' % len(val_dataset))
############################################################
net = network((3,640),num_classes = num_class,resnet_factory = models.resnet101, freeze_resnet=False)
start_epoch = 1
lr = base_lr
best_val_loss = float('inf')
log_mode = 'w'
if os.path.exists(args.resume):
print('loading checkpoint %s'%(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch'] + 1
lr = checkpoint['lr']
best_val_loss = checkpoint['best_val_loss']
net.load_state_dict(checkpoint['state_dict'])
log_mode = 'a'
net = net.cuda()
parser.add_argument('-d', '--dataset', default='NYU', help='NYU or SUN', type=str)
args = parser.parse_args()
save_dir = './%s_RDFnet/'%args.dataset
if not os.path.exists(save_dir):
os.mkdir(save_dir)
val_dataset = RGBD(args.dataset,'val')
val_loader = DataLoader(
val_dataset, batch_size=batch_size, shuffle=False, num_workers=workers)
num_class = {'SUN':37,'NYU':40}[args.dataset]
ignore_label = {'SUN':255,'NYU':-1}[args.dataset]
loss = nn.CrossEntropyLoss(ignore_index=ignore_label)
patience = {'SUN':10,'NYU':40}[args.dataset]
print('Val sample number: %d' % len(val_dataset))
############################################################
net = network(640,num_classes = num_class,resnet_factory = models.resnet152, freeze_resnet=False)
loss = nn.CrossEntropyLoss(ignore_index=ignore_label)
net = net.cuda()
loss = loss.cuda()
cudnn.benchmark = True
net = DataParallel(net)
print('loading checkpoint ')
checkpoint = torch.load(val_path)
net.load_state_dict(checkpoint['state_dict'])