()
| 128 | |
| 129 | |
| 130 | def main(): |
| 131 | args = parse_args() |
| 132 | # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule? |
| 133 | args.amodel = args.amodel.replace("/", "-") |
| 134 | # download sizes.json file |
| 135 | |
| 136 | # (yusong): the below two lines are for debug |
| 137 | # print("setting up faulthandler") |
| 138 | # faulthandler.register(10) |
| 139 | |
| 140 | random.seed(args.seed) |
| 141 | torch.manual_seed(args.seed) |
| 142 | torch.cuda.manual_seed(args.seed) |
| 143 | torch.cuda.manual_seed_all(args.seed) |
| 144 | np.random.seed(args.seed) |
| 145 | if args.tmodel == "bert" or args.tmodel == "roberta" or args.tmodel == "bart": |
| 146 | assert ( |
| 147 | args.pretrained == "" or args.pretrained is None |
| 148 | ), "bert/roberta/bart text encoder does not support pretrained models." |
| 149 | |
| 150 | # get the name of the experiments |
| 151 | if args.name is None: |
| 152 | args.name = "-".join( |
| 153 | [ |
| 154 | datetime.now().strftime("%Y_%m_%d-%H_%M_%S"), |
| 155 | f"model_{args.amodel}", |
| 156 | f"lr_{args.lr}", |
| 157 | f"b_{args.batch_size}", |
| 158 | f"j_{args.workers}", |
| 159 | f"p_{args.precision}", |
| 160 | ] |
| 161 | ) |
| 162 | |
| 163 | # discover initial world args early so we can log properly |
| 164 | args.distributed = False |
| 165 | args.local_rank, args.rank, args.world_size = world_info_from_env() |
| 166 | |
| 167 | if args.remotedata and is_master(args): |
| 168 | for dataset_name in args.datasetnames: |
| 169 | for split in dataset_split[dataset_name]: |
| 170 | if not os.path.exists(f"./json_files/{dataset_name}/{split}"): |
| 171 | os.makedirs(f"./json_files/{dataset_name}/{split}") |
| 172 | os.system( |
| 173 | f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json" |
| 174 | ) |
| 175 | |
| 176 | args.log_path = None |
| 177 | if is_master(args, local=args.log_local): |
| 178 | log_base_path = os.path.join(args.logs, args.name) |
| 179 | os.makedirs(log_base_path, exist_ok=True) |
| 180 | log_filename = f"out-{args.rank}" if args.log_local else "out.log" |
| 181 | args.log_path = os.path.join(log_base_path, log_filename) |
| 182 | if os.path.exists(args.log_path): |
| 183 | print( |
| 184 | "Error. Experiment already exists. Use --name {} to specify a new experiment." |
| 185 | ) |
| 186 | return -1 |
| 187 |
no test coverage detected