speech_net = torch.nn.DataParallel(speech_net).cuda()
# define optimizer
if args.opt.lower() == 'adam':
optimizer = optim.Adam(speech_net.parameters(), lr=args.lr)
elif args.opt.lower() == 'sgd':
optimizer = optim.SGD(speech_net.parameters(),
lr=args.lr,
momentum=args.momentum)
else:
optimizer = optim.SGD(speech_net.parameters(),
lr=args.lr,
momentum=args.momentum)
train_dataset = Datasets.SpeechYoloDataSet(classes_root_dir=args.train_data,
this_root_dir=args.train_data,
yolo_config=config_dict,
augment=args.augment_data)
val_dataset = Datasets.SpeechYoloDataSet(classes_root_dir=args.train_data,
this_root_dir=args.val_data,
yolo_config=config_dict)
sampler_train = Datasets.ImbalancedDatasetSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=20,
pin_memory=args.cuda,
sampler=sampler_train)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
args = parser.parse_args()
args.cuda = args.cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
model, acc, epoch = load_model(args.model)
config_dict = {"C": model.c, "B": model.b, "K": model.k}
if args.cuda:
print('Using CUDA with {0} GPUs'.format(torch.cuda.device_count()))
model = torch.nn.DataParallel(model).cuda()
test_dataset = Datasets.SpeechYoloDataSet(classes_root_dir=args.train_data,
this_root_dir=args.test_data,
yolo_config=config_dict,
words_list_file=args.word_list)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=20,
pin_memory=args.cuda,
sampler=None)
range_split = args.theta_range.split('_')
start_theta = float(range_split[0])
end_theta = float(range_split[1])
step_theta = float(range_split[2])
loss = loss_speech_yolo.YOLOLoss()
# threshold = args.decision_threshold
args = parser.parse_args()
args.cuda = args.cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
model, acc, epoch = load_model(args.model)
config_dict = {"C": model.c, "B": model.b, "K": model.k}
if args.cuda:
print('Using CUDA with {0} GPUs'.format(torch.cuda.device_count()))
model = torch.nn.DataParallel(model).cuda()
test_dataset = Datasets.SpeechYoloDataSet(classes_root_dir=args.train_data,
this_root_dir=args.test_data,
yolo_config=config_dict)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=20,
pin_memory=args.cuda,
sampler=None)
range_split = args.theta_range.split('_')
start_theta = float(range_split[0])
end_theta = float(range_split[1])
step_theta = float(range_split[2])
loss = loss_speech_yolo.YOLOLoss()
# threshold = args.decision_threshold
