def main():
train_dataset = hrb_input.TrainDataset()
val_dataset = hrb_input.ValidationDataset()
test_dataset = hrb_input.TestDataset()
train_loader = DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8,
pin_memory=True)
val_loader = DataLoader(val_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8,
pin_memory=True)
test_loader = DataLoader(test_dataset,
batch_size=BATCH_SIZE,
num_workers=8,
pin_memory=True)
model = network.resnet50(num_classes=1000)
model = DataParallel(model, device_ids=[0, 1, 2])
model.to(device)
loss_func = F.cross_entropy
optimizer = optim.Adam(model.parameters(), lr=LR)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=LR_STEPS,
gamma=LR_GAMMA)
best_loss = 100.0
patience_counter = 0
for epoch in range(EPOCHS):
train(model, optimizer, loss_func, train_loader, epoch)
val_loss = validate(model, loss_func, val_loader)
if val_loss < best_loss:
torch.save(model, MODEL_PATH)
print('Saving improved model')
print()
best_loss = val_loss
patience_counter = 0
else:
patience_counter += 1
print('Epoch(s) since best model: ', patience_counter)
print()
if patience_counter >= EARLY_STOPPING_EPOCHS:
print('Early Stopping ...')
print()
break
scheduler.step()
print('Predicting labels from best trained model')
predict(test_loader)
def train_model(model, model_name, hyperparams, device, epochs):
'''
Train Model
This is a generic function to call the model's training function.
'''
print('Beginning Training for: ', model_name)
print('------------------------------------')
results = {}
if torch.cuda.device_count() > 1:
print("Using ", torch.cuda.device_count(), " GPUs.")
print('------------------------------------')
model = DataParallel(model)
model = model.to(device=device)
optimizer = optim.Adam(model.parameters(), betas=hyperparams['betas'], lr=hyperparams['learning_rate'], weight_decay=hyperparams['L2_reg'])
lr_updater = lr_scheduler.StepLR(optimizer, hyperparams['lr_decay_epochs'], hyperparams['lr_decay'])
results = train(model, optimizer, lr_updater, results, epochs=epochs)
plot_results(results, model_name ,save=True)
np.save(model_name, results)
return results
def set_device(self, device):
device = cast_device(device)
str_device = device_to_str(device)
nn_module = self.get_nn_module()
if isinstance(device, (list, tuple)):
device_ids = []
for dev in device:
if dev.type != 'cuda':
raise ValueError
if dev.index is None:
raise ValueError
device_ids.append(dev.index)
if len(device_ids) != len(set(device_ids)):
raise ValueError("Cuda device indices must be unique")
nn_module = DataParallel(nn_module, device_ids=device_ids)
device = device[0]
self.params['device'] = str_device
self.device = device
self.nn_module = nn_module.to(self.device)
if self.loss is not default:
self.loss = self.loss.to(self.device)
def neuralwarp_train(**kwargs):
# 多尺度图片训练 396+
print(kwargs)
#print("Mask == 1")
with open(kwargs['params']) as f:
params = json.load(f)
if kwargs['manner'] == 'train':
params['is_train'] = True
else:
params['is_train'] = False
params['batch_size'] = kwargs['batch_size']
if torch.cuda.device_count() > 1:
print("-------------------Parallel_GPU_Train--------------------------")
parallel = True
else:
print("------------------Single_GPU_Train----------------------")
parallel = False
opt.feature = 'cqt'
opt.notes = 'SoftDTW'
opt.model = 'SoftDTW'
opt.batch_size = 'batch_size'
os.environ["CUDA_VISIBLE_DEVICES"] = str(kwargs["Device"])
opt.Device=kwargs["Device"]
#device_ids = [2]
opt._parse(kwargs)
model = getattr(models, opt.model)(params)
p = 'check_points/' + model.model_name + opt.notes
#f = os.path.join(p, "0620_07:05:30.pth")#使用Neural_dtw目前最优 0620_07:05:30.pth cover80 map:0.705113267654046 0.08125 7.96875
#f = os.path.join(p, "0620_17:37:35.pth")
#f = os.path.join(p, "0621_22:42:59.pth")#NeuralDTW_Milti_Metix_res 0622_16:33:07.pth 0621_22:42:59.pth
#f = os.path.join(p, "0628_17:00:52.pth")#0628_17:00:52.pth FCN
#f = os.path.join(p,"0623_16:01:05.pth") #3seq
#f = os.path.join(p,"0630_07:59:56.pth")#VGG11 0630_01:10:15.pth 0630_07:59:56.pth
if kwargs['model'] == 'NeuralDTW_CNN_Mask_dilation_SPP':
f = os.path.join(p,"0704_19:58:25.pth")
elif kwargs['model'] == 'NeuralDTW_CNN_Mask_dilation_SPP2':
f = os.path.join(p,"0709_00:31:23.pth")
elif kwargs['model'] == 'NeuralDTW_CNN_Mask_dilation':
f = os.path.join(p,"0704_06:40:41.pth")
opt.load_model_path = f
if kwargs['model'] != 'NeuralDTW' and kwargs['manner'] != 'train':
if opt.load_latest is True:
model.load_latest(opt.notes)
elif opt.load_model_path:
print("load_model:",opt.load_model_path)
model.load(opt.load_model_path)
if parallel == True:
model = DataParallel(model)
model.to(opt.device)
torch.multiprocessing.set_sharing_strategy('file_system')
# step2: data
out_length =400
if kwargs['model'] == 'NeuralDTW_CNN_Mask_300':
out_length = 300
if kwargs['model'] == 'NeuralDTW_CNN_Mask_spp':
train_data0 = triplet_CQT(out_length=200, is_label=kwargs['is_label'], is_random=kwargs['is_random'])
train_data1 = triplet_CQT(out_length=300, is_label=kwargs['is_label'], is_random=kwargs['is_random'])
train_data2 = triplet_CQT(out_length=400, is_label=kwargs['is_label'], is_random=kwargs['is_random'])
else:
train_data0 = triplet_CQT(out_length=out_length, is_label=kwargs['is_label'], is_random=kwargs['is_random'])
train_data1 = triplet_CQT(out_length=out_length, is_label=kwargs['is_label'], is_random=kwargs['is_random'])
train_data2 = triplet_CQT(out_length=out_length, is_label=kwargs['is_label'], is_random=kwargs['is_random'])
val_data80 = CQT('songs80', out_length=kwargs['test_length'])
val_data = CQT('songs350', out_length=kwargs['test_length'])
val_data_marukars = CQT('Mazurkas',out_length=kwargs['test_length'])
train_dataloader0 = DataLoader(train_data0, opt.batch_size, shuffle=True, num_workers=opt.num_workers)
train_dataloader1 = DataLoader(train_data1, opt.batch_size, shuffle=True, num_workers=opt.num_workers)
train_dataloader2 = DataLoader(train_data2, opt.batch_size, shuffle=True, num_workers=opt.num_workers)
val_dataloader80 = DataLoader(val_data80, 1, shuffle=False, num_workers=1)
val_dataloader = DataLoader(val_data, 1, shuffle=False, num_workers=1)
val_dataloader_marukars = DataLoader(val_data_marukars,1, shuffle=False, num_workers=1)
if kwargs['manner'] == 'test':
# val_slow(model, val_dataloader, style='null')
val_slow_batch(model,val_dataloader_marukars, batch=100, is_dis=kwargs['zo'])
elif kwargs['manner'] == 'visualize':
visualize(model, val_dataloader80)
elif kwargs['manner'] == 'mul_test':
p = 'check_points/' + model.model_name + opt.notes
l = sorted(os.listdir(p))[: 20]
best_MAP, MAP = 0, 0
for f in l:
f = os.path.join(p, f)
model.load(f)
model.to(opt.device)
MAP += val_slow_batch(model, val_dataloader, batch=400, is_dis=kwargs['zo'])
MAP += val_slow_batch(model, val_dataloader80, batch=400, is_dis=kwargs['zo'])
if MAP > best_MAP:
print('--best result--')
best_MAP = MAP
MAP = 0
else:
# step3: criterion and optimizer
be = torch.nn.BCELoss()
lr = opt.lr
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=opt.weight_decay)
# if parallel is True:
# optimizer = torch.optim.Adam(model.module.parameters(), lr=lr, weight_decay=opt.weight_decay)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.8, patience=10, verbose=True, min_lr=5e-6)
# step4: train
best_MAP = 0
for epoch in range(opt.max_epoch):
running_loss = 0
num = 0
for ii, ((a0, p0, n0, la0, lp0, ln0), (a1, p1, n1, la1, lp1, ln1), (a2, p2, n2, la2, lp2, ln2)) in tqdm(
enumerate(zip(train_dataloader0, train_dataloader1, train_dataloader2))):
# for ii, (a2, p2, n2) in tqdm(enumerate(train_dataloader2)):
for flag in range(3):
if flag == 0:
a, p, n, la, lp, ln = a0, p0, n0, la0, lp0, ln0
elif flag == 1:
a, p, n, la, lp, ln = a1, p1, n1, la1, lp1, ln1
else:
a, p, n, la, lp, ln = a2, p2, n2, la2, lp2, ln2
B, _, _, _ = a.shape
if kwargs["zo"] == True:
target = torch.cat((torch.zeros(B), torch.ones(B))).cuda()
else:
target = torch.cat((torch.ones(B), torch.zeros(B))).cuda()
# train model
a = a.requires_grad_().to(opt.device)
p = p.requires_grad_().to(opt.device)
n = n.requires_grad_().to(opt.device)
optimizer.zero_grad()
pred = model(a, p, n)
pred = pred.squeeze(1)
loss = be(pred, target)
loss.backward()
optimizer.step()
running_loss += loss.item()
num += a.shape[0]
if ii % 5000 == 0:
running_loss /= num
print("train_loss:",running_loss)
MAP = 0
print("Youtube350:")
MAP += val_slow_batch(model, val_dataloader, batch=1 , is_dis=kwargs['zo'])
print("CoverSong80:")
MAP += val_slow_batch(model, val_dataloader80, batch=1, is_dis=kwargs['zo'])
# print("Marukars:")
# MAP += val_slow_batch(model, val_dataloader_marukars, batch=100, is_dis=kwargs['zo'])
if MAP > best_MAP:
best_MAP = MAP
print('*****************BEST*****************')
if kwargs['save_model'] == True:
if parallel:
model.module.save(opt.notes)
else:
model.save(opt.notes)
scheduler.step(running_loss)
running_loss = 0
num = 0