val_sampler = torch.utils.data.distributed.DistributedSampler( val_dataset, num_replicas=hvd.size(), rank=hvd.rank()) val_loader = DataLoader(val_dataset, batch_size=val_batch_size, collate_fn=val_dataset.collate, sampler=val_sampler, **kwargs) # Horovod: print logs on the first worker. verbose = 1 if hvd.rank() == 0 else 0 # ------------ preparation ------------ net = SCNN(resize_shape, pretrained=True) lr_scaler = 1 if torch.cuda.is_available(): net.cuda() # Horovod: Scale learning rate as per number of devices if hvd.nccl_built(): lr_scaler = hvd.local_size() net = torch.nn.DataParallel(net) lr = exp_cfg['optim']['lr'] momentum = exp_cfg['optim']['momentum'] weight_decay = exp_cfg['optim']['weight_decay'] nesterov = exp_cfg['optim']['nesterov'] # Horovod: scale learning rate by lr_scaler. optimizer = optim.SGD(net.parameters(), lr=lr * lr_scaler, momentum=momentum, weight_decay=weight_decay,
import cv2 import torch from model import SCNN from utils.prob2lines import getLane from utils.transforms import * import time from tqdm import tqdm import pdb image_resize_width = 512 net = SCNN(input_size=(image_resize_width, 288), pretrained=False) if torch.cuda.is_available(): net = net.cuda() mean=(0.3598, 0.3653, 0.3662) # CULane mean, std std=(0.2573, 0.2663, 0.2756) transform = Compose(Resize((image_resize_width, 288)), ToTensor(), Normalize(mean=mean, std=std)) def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--img_path", '-i', type=str, default="demo/demo.jpg", help="Path to demo img") parser.add_argument("--weight_path", '-w', type=str, help="Path to model weights") parser.add_argument("--visualize", '-v', action="store_true", default=False, help="Visualize the result") args = parser.parse_args() return args