parser.add_argument('--gpu_ids', type=str, default='0', help='use which gpu to train, must be a \ comma-separated list of integers only (default=0)') args = parser.parse_args() args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device('cuda' if args.cuda else 'cpu') if args.cuda: try: args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')] except ValueError: raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only') torch.cuda.manual_seed(args.seed) torch.manual_seed(args.seed) print(args) train_set = get_training_set(args.upscale_factor) test_set = get_test_set(args.upscale_factor) training_data_loader = DataLoader(dataset=train_set, num_workers=args.threads, batch_size=args.batch_size, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=args.threads, batch_size=args.test_batch_size, shuffle=False) srcnn = SRCNN() criterion = nn.MSELoss() if args.cuda: srcnn = nn.DataParallel(srcnn, device_ids=args.gpu_ids) srcnn = srcnn.to(device) criterion = criterion.to(device) parameters = [ conv1_rgb
type=int, default=123, help='random seed to use. Default=123') opt = parser.parse_args() print(opt) use_cuda = opt.cuda if use_cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) if use_cuda: torch.cuda.manual_seed(opt.seed) train_set = get_training_set(opt.upscale_factor) test_set = get_test_set(opt.upscale_factor) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.test_batch_size, shuffle=False) srcnn = SRCNN() criterion = nn.MSELoss() if (use_cuda): srcnn.cuda()
parser = argparse.ArgumentParser( description='This program quantizes weight by using weight sharing') parser.add_argument('model', type=str, help='path to saved pruned model') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--output', default='saves/model_after_weight_sharing_6bits.ptmodel', type=str, help='path to model output') args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() model = torch.load(args.model) train_set = get_training_set(2) test_set = get_test_set(2) training_data_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=50, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=4, batch_size=10, shuffle=False) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.001) initial_optimizer_state_dict = optimizer.state_dict()