type=bool,
default=True,
metavar='B',
help='mask out known values (default: True)')
parser.add_argument('--colored',
type=bool,
default=True,
metavar='B',
help='uses model with color information (default: True)')
# SETUP
print('SETUP')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
args.device = torch.device('cuda:0') if use_cuda else torch.device('cpu')
saveDir = os.path.join('../models/', args.expID)
writeArgsFile(args, saveDir)
torch.manual_seed(args.seed)
kwargs = {}
print('Seed: {:d}'.format(args.seed))
print('Device: {}'.format(args.device))
if use_cuda:
print('\nCUDA')
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.benchmark = args.benchmark
num_gpus = torch.cuda.device_count()
kwargs = {'num_workers': num_gpus * args.workers, 'pin_memory': True}
print("Number of GPUs: {:d}".format(num_gpus))
print('Workers/GPU: {:d}'.format(args.workers))
if args.resume and not args.model:
print("\n=> No model to resume training. Double-check arguments!")
quit()
torch.manual_seed(args.seed)
kwargs = {}
if args.cuda:
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.benchmark = True
kwargs = {'num_workers': 1, 'pin_memory': True}
if not os.path.exists(args.saveDir):
os.makedirs(args.saveDir)
utils.writeArgsFile(args, args.saveDir)
## LOAD DATASETS
print('\nDATASET INFO.')
train_data = TrainDataHandler('../datasets/train/')
test_data = TestDataHandler('../datasets/test/')
print('Train & val. size: {} x {}'.format(len(train_data),
train_data[0][0].size()))
print('Test size: {} x {}'.format(len(test_data), test_data[0][0].size()))
## LOAD MODEL & SOLVER
print('\nLOADING NETWORK & SOLVER.')
model = MyNet() if not args.vnect else VNect()