args.model_name,
args.backbone, cur_time)
# writer = SummaryWriter(args.record_folder)
os.makedirs(args.record_folder, exist_ok=True)
if args.use_cuda and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
elif args.use_cuda:
raise RuntimeError(
'You are using GPU mode, but GPUs are not available!')
# construct model and optimizer
assert (args.train_data == args.test_data)
args.image_len = 84
args.semantic_size = get_semantic_size(args)
args.out_channels, args.feature_h = get_outputs_c_h(
args.backbone, args.image_len)
model = ProtoNetAGAMwoAttr(args.backbone, args.semantic_size,
args.out_channels)
if args.use_cuda:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
num_gpus = torch.cuda.device_count()
if args.multi_gpu:
model = nn.DataParallel(model)
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=0.0005)
args = parser.parse_args()
# make folder and tensorboard writer to save model and results
cur_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
args.record_folder = './{}_{}_{}_{}-ways_{}-shots_{}'.format(args.train_data, args.test_data, args.backbone, str(args.num_ways), str(args.num_shots), cur_time)
os.makedirs(args.record_folder, exist_ok=True)
if args.use_cuda and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
elif args.use_cuda:
raise RuntimeError('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)
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)