（17-7-01）基于强化学习的自动驾驶系统：生成训练数据+训练模型

 from argparse import ArgumentParser

    
  
    
 from data.collect_data import main
    
 from configuration.config import *
    
  
    
 if __name__=='__main__':
    
     argparser = ArgumentParser()
    
     argparser.add_argument('--world-port', type=int, default=WORLD_PORT)
    
     argparser.add_argument('--host', type=str, default=WORLD_HOST)
    
     argparser.add_argument('--map', type=str, default=TRAIN_MAP, help="Load the map before starting the simulation")
    
     argparser.add_argument('--weather', type=str, default='ClearNoon',
    
                        choices=['ClearNoon', 'ClearSunset', 'CloudyNoon', 'CloudySunset',
    
                                 'WetNoon', 'WetSunset', 'MidRainyNoon', 'MidRainSunset',
    
                                 'WetCloudyNoon', 'WetCloudySunset', 'HardRainNoon',
    
                                 'HardRainSunset', 'SoftRainNoon', 'SoftRainSunset'],
    
                                 help='Weather preset')
    
     argparser.add_argument('--out_folder', type=str, default='./sensor_data', help="Output folder")
    
     argparser.add_argument('--nb_frames', type=int, default=300, help="Number of frames to record per route")
    
     argparser.add_argument('--nb_passes', type=int, default=7, help="Number of passes per route")
    
     argparser.add_argument('--freq_save', type=int, default=5, help="Frequency of saving data (in steps)")
    
     argparser.add_argument('-np', action='store_true', help='Save data as numpy arrays instead of images')
    
     argparser.add_argument('--begin', type=int, default=0, help='Begin at this episode (for resuming)')
    
  
    
     args = argparser.parse_args()
    
  
    
     main(args)

 def main(args):

    
     os.makedirs(args.dirpath, exist_ok=True)
    
  
    
     img_size = [int(x) for x in args.img_size.split('x')] if args.img_size != 'default' else None
    
  
    
     if args.split:
    
     with open(args.split, 'rb') as f:
    
         split = pickle.load(f)
    
     train = split['train']
    
     val = split['val']
    
     
    
     else:
    
     with open(args.file, 'rb') as f:
    
         data = pickle.load(f)
    
  
    
     stratify = [d[0] for d in data]
    
  
    
     # Split data stratify by folder (different weather conditions)
    
     train, val = train_test_split(data, test_size=args.val_size, random_state=42,
    
                                 shuffle=True, stratify=stratify)
    
     
    
     #Save split
    
     with open(args.dirpath + '/split.pkl', 'wb') as f:
    
         pickle.dump({'train': train, 'val': val}, f)
    
  
    
     
    
     normalize_output = False if args.model == 'AutoencoderSEM' else True
    
     train_dataset = AutoencoderDataset(train[:100], img_size, args.norm_input, args.low_sem, args.use_img_out, normalize_output)
    
     val_dataset = AutoencoderDataset(val[:100], img_size, args.norm_input, args.low_sem, args.use_img_out, normalize_output)
    
  
    
     train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    
     val_loader = DataLoader(val_dataset, batch_size=args.batch_size*2, shuffle=False)
    
  
    
     img_size = tuple(train_dataset[0][0].shape[1:]) if img_size is None else img_size
    
     num_classes = 14 if args.low_sem else 29
    
     
    
     if args.model == 'Autoencoder':
    
     model = Autoencoder(input_size=img_size, emb_size=args.emb_size, lr=args.lr, weights=(0.8, 0.1, 0.1),
    
                         use_additional_data=args.additional_data, out_ch=4 if args.use_img_out else 1)
    
     elif args.model == 'AutoencoderSEM':
    
     model = AutoencoderSEM(input_size=img_size, emb_size=args.emb_size, lr=args.lr, weights=(0.8, 0.1, 0.1),
    
                         use_additional_data=args.additional_data, num_classes=num_classes)
    
     elif args.model == 'VAE':
    
     model = VAE(input_size=img_size, emb_size=args.emb_size, lr=args.lr,
    
                         out_ch=4 if args.use_img_out else 1)
    
     else:
    
     raise ValueError('Model not found')
    
  
    
     trainer = pl.Trainer(callbacks=[pl.callbacks.ModelCheckpoint(dirpath=args.dirpath, monitor='val_loss', save_top_k=1),
    
                                 pl.callbacks.LearningRateMonitor(logging_interval='epoch'),
    
                                 pl.callbacks.ModelCheckpoint(dirpath=args.dirpath, filename="{epoch}")],
    
                      accelerator=args.device,
    
                      max_epochs=args.epochs,
    
                      default_root_dir=args.dirpath)
    
     trainer.fit(model, train_loader, val_loader, ckpt_path=args.pretrained)
    
     
    
  
    
 if __name__=='__main__':
    
     parser = ArgumentParser()
    
     parser.add_argument('--file', type=str, default=AE_DATASET_FILE, help='dataset file')
    
     parser.add_argument('--val_size', type=float, default=AE_VAL_SIZE, help='validation size')
    
     parser.add_argument('--img_size', type=str, default=AE_IMG_SIZE, help='image size')
    
  
    
     parser.add_argument('--norm_input', type=bool, default=AE_NORM_INPUT, help='Normalize input image')
    
     parser.add_argument('--emb_size', type=int, default=AE_EMB_SIZE, help='embedding size')
    
     parser.add_argument('--batch_size', type=int, default=AE_BATCH_SIZE, help='batch size')
    
     parser.add_argument('--epochs', type=int, default=AE_EPOCHS, help='number of epochs')
    
     parser.add_argument('--lr', type=float, default=AE_LR, help='learning rate')
    
     parser.add_argument('--device', type=str, default='auto', help='device', choices=['auto', 'gpu', 'cpu'])
    
     parser.add_argument('--dirpath', type=str, default=AE_DIRPATH, help='directory path to save the model')
    
  
    
     parser.add_argument('--low_sem', type=bool, default=AE_LOW_SEM, help='Use low resolution semantic segmentation (14 classes)')
    
     parser.add_argument('--use_img_out', type=bool, default=AE_USE_IMG_AS_OUTPUT, help='Use image as output (not used if model is VAE)')
    
     parser.add_argument('--model', type=str, default=AE_MODEL, help='model', choices=['Autoencoder', 'AutoencoderSEM', 'VAE'])
    
     parser.add_argument('--additional_data', type=bool, default=AE_ADDITIONAL_DATA, help='Use additional data (not used if model is VAE)')
    
  
    
     parser.add_argument('--pretrained', type=str, default=AE_PRETRAINED, help='pretrained model')
    
     parser.add_argument('--split', type=str, default=AE_SPLIT, help='split file')
    
     args = parser.parse_args()
    
  
    
     main(args)