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()