def main():
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('data', metavar='DIR', help='path to dataset')
parser.add_argument('--batch-size',
type=int,
default=128,
metavar='N',
help='input batch size for training (default: 128)')
parser.add_argument('--batch-size-val',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--epochs',
type=int,
default=1000,
metavar='N',
help='number of epochs to train (default: 1000)')
parser.add_argument('--lr',
type=float,
default=1e-4,
metavar='LR',
help='learning rate (default: 1e-4)')
parser.add_argument('--image-size',
type=float,
default=80,
metavar='IMSIZE',
help='input image size (default: 80)')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--multi-gpu',
action='store_true',
default=False,
help='parallel training on multiple GPUs')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument('--save-model',
action='store_true',
default=False,
help='For Saving the current Model')
parser.add_argument('--model-save-path',
default='',
type=str,
metavar='PATH',
help='For Saving the current Model')
parser.add_argument('--resume',
default='',
type=str,
metavar='PATH',
help='path to latest checkpoint (default: none)')
args = parser.parse_args()
torch.set_default_tensor_type('torch.FloatTensor')
device = torch.device("cpu" if args.no_cuda else "cuda")
train_data = dataset(args.data,
"train",
args.image_size,
transform=transforms.Compose([ToTensor()]),
shuffle=True)
valid_data = dataset(args.data,
"val",
args.image_size,
transform=transforms.Compose([ToTensor()]))
trainloader = DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=8)
validloader = DataLoader(valid_data,
batch_size=args.batch_size_val,
shuffle=False,
num_workers=8)
model = Model(args.image_size, args.image_size)
optimizer = RAdam(model.parameters(), lr=args.lr, weight_decay=1e-8)
if not args.no_cuda:
model.cuda()
if torch.cuda.device_count() > 1 and args.multi_gpu:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = torch.nn.DataParallel(model)
if args.resume:
model.load_state_dict(
torch.load(os.path.join(args.resume, model_save_name)))
optimizer.load_state_dict(
torch.load(os.path.join(args.resume, optimizer_save_name)))
train(model, optimizer, trainloader, validloader, device, args)
def main():
parser = argparse.ArgumentParser(description='RAVEN test args')
parser.add_argument('data', metavar='DIR', help='path to dataset')
parser.add_argument('--batch-size',
type=int,
default=100,
metavar='N',
help='input batch size for training (default: 100)')
parser.add_argument('--image-size',
type=float,
default=80,
metavar='IMSIZE',
help='input image size (default: 80)')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--multi-gpu',
action='store_true',
default=False,
help='parallel training on multiple GPUs')
parser.add_argument('--valid_result',
action='store_true',
default=False,
help='compute results on validation dataset')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument('--model-path-name',
default='',
type=str,
metavar='PATH',
help='The path+name of model to be loaded')
args = parser.parse_args()
torch.set_default_tensor_type('torch.FloatTensor')
device = torch.device("cpu" if args.no_cuda else "cuda")
test_data = dataset(args.data,
"test",
args.image_size,
transform=transforms.Compose([ToTensor()]))
valid_data = dataset(args.data,
"val",
args.image_size,
transform=transforms.Compose([ToTensor()]))
testloader = DataLoader(test_data,
batch_size=args.batch_size,
num_workers=8)
validloader = DataLoader(valid_data,
batch_size=args.batch_size,
num_workers=8)
model = Model(args.image_size, args.image_size)
if not args.no_cuda:
model.cuda()
if torch.cuda.device_count() > 1 and args.multi_gpu:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load(args.model_path_name))
test(model, testloader, validloader, device, args)