parser.add_argument('--opt', type=str, default='adam', help='') parser.add_argument('--lr', type=float, default=0.001, help='') parser.add_argument('--weight_decay', type=float, default=0, help='') parser.add_argument('--early_stopping', type=int, default=20, help='') parser.add_argument('--save_epochs', type=str, default='5,10,15,20,25', help='') parser.add_argument('--save_every_epoch', type=int, default=26, help='') #26(for MovieLens), 16(only for Yelp) args = parser.parse_args() # Setup data and weights file path data_folder, weights_folder, logger_folder = \ get_folder_path(model=MODEL, dataset=args.dataset + args.dataset_name, loss_type=LOSS_TYPE) # Setup device if not torch.cuda.is_available() or args.device == 'cpu': device = 'cpu' else: device = 'cuda:{}'.format(args.gpu_idx) # Setup args dataset_args = { 'root': data_folder, 'dataset': args.dataset, 'name': args.dataset_name, 'if_use_features': args.if_use_features.lower() == 'true', 'num_negative_samples': args.num_negative_samples, 'num_core': args.num_core,
parser.add_argument("--batch_size", type=int, default=128, help="") parser.add_argument("--num_workers", type=int, default=16, help="") parser.add_argument("--opt", type=str, default='adam', help="") parser.add_argument("--loss", type=str, default='mse', help="") parser.add_argument("--lr", type=float, default=1e-3, help="") parser.add_argument("--weight_decay", type=float, default=1e-3, help="") parser.add_argument("--early_stopping", type=int, default=60, help="") parser.add_argument("--save_epochs", type=list, default=[10, 40, 80], help="") parser.add_argument("--save_every_epoch", type=int, default=40, help="") args = parser.parse_args() # Setup data and weights file path data_folder, weights_folder, logger_folder = \ get_folder_path(model=MODEL, dataset=args.dataset + args.dataset_name) # Setup device if not torch.cuda.is_available() or args.device == 'cpu': device = 'cpu' else: device = 'cuda:{}'.format(args.gpu_idx) # Setup args dataset_args = { 'root': data_folder, 'dataset': args.dataset, 'name': args.dataset_name, 'if_use_features': args.if_use_features, 'num_core': args.num_core, 'num_feat_core': args.num_feat_core, 'train_ratio': args.train_ratio } model_args = {