# construct model and optimizer
args.image_len = 28 if args.train_data == 'omniglot' else 84
args.out_channels, _ = get_outputs_c_h(args.backbone, args.image_len)
model = ConvolutionalNeuralNetwork(args.backbone, args.out_channels, args.num_ways)
teacher_model = ConvolutionalNeuralNetwork(args.teacher_backbone, args.out_channels, args.num_ways)
if args.use_cuda:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
num_gpus = torch.cuda.device_count()
if args.multi_gpu:
model = DataParallel(model)
teacher_model = DataParallel(teacher_model)
model = model.cuda()
teacher_model = teacher_model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=0.0005)
# download teacher trained model
if args.teacher_resume and args.teacher_resume_folder is not None:
teacher_model_path = os.path.join(args.teacher_resume_folder, ('_'.join(['teacher', args.train_data, args.test_data, args.teacher_backbone, 'max_acc']) + '.pt'))
teacher_state = torch.load(teacher_model_path)
if args.multi_gpu:
teacher_model.module.load_state_dict(teacher_state)
else:
teacher_model.load_state_dict(teacher_state)
# training from the checkpoint
if args.resume and args.resume_folder is not None:
'You are using GPU mode, but GPUs are not available!')
# construct model and optimizer
args.image_len = 28 if args.train_data == 'omniglot' else 84
args.out_channels, _ = get_outputs_c_h(args.backbone, args.image_len)
model = ConvolutionalNeuralNetwork(args.backbone, args.out_channels,
args.num_ways)
if args.use_cuda:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
num_gpus = torch.cuda.device_count()
if args.multi_gpu:
model = DataParallel(model)
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=0.0005)
# training from the checkpoint
if args.resume and args.resume_folder is not None:
# load checkpoint
checkpoint_path = os.path.join(args.resume_folder, ('_'.join([
args.model_name, args.train_data, args.test_data, args.backbone,
'max_acc'
]) + '_checkpoint.pt.tar')) # tag='max_acc' can be changed
state = torch.load(checkpoint_path)
if args.multi_gpu:
model.module.load_state_dict(state['model'])