class Test:
def __init__(self, model_name, snapshot, num_classes):
self.num_classes = num_classes
if model_name == "resnet50":
self.model = ResNet(torchvision.models.resnet50(pretrained=False),
num_classes)
elif model_name == "mobilenetv2":
self.model = MobileNetV2(num_classes=num_classes)
else:
print("No such model...")
exit(0)
self.saved_state_dict = torch.load(snapshot)
self.model.load_state_dict(self.saved_state_dict)
self.model.cuda(0)
self.model.eval() # Change model to 'eval' mode
self.softmax = nn.Softmax(dim=1).cuda(0)
self.optimizer = Optimize()
def draw_vectors(self, pred_vector1, pred_vector2, pred_vector3, img,
center, width):
optimize_v = self.optimizer.Get_Ortho_Vectors(np.array(pred_vector1),
np.array(pred_vector2),
np.array(pred_vector3))
v1, v2, v3 = optimize_v[0], optimize_v[1], optimize_v[2]
# draw vector in blue color
predx, predy, predz = v1
utils.draw_front(img,
predx,
predy,
width,
tdx=center[0],
tdy=center[1],
size=100,
color=(255, 0, 0))
# draw vector in green color
predx, predy, predz = v2
utils.draw_front(img,
predx,
predy,
width,
tdx=center[0],
tdy=center[1],
size=100,
color=(0, 255, 0))
# draw vector in red color
predx, predy, predz = v3
utils.draw_front(img,
predx,
predy,
width,
tdx=center[0],
tdy=center[1],
size=100,
color=(0, 0, 255))
cv.imshow("pose visualization", img)
def test_per_img(self, cv_img, draw_img, center, w):
with torch.no_grad():
images = cv_img.cuda(0)
# get x,y,z cls predictions
x_v1, y_v1, z_v1, x_v2, y_v2, z_v2, x_v3, y_v3, z_v3 = self.model(
images)
# get prediction vector(get continue value from classify result)
_, _, _, pred_vector1 = utils.classify2vector(
x_v1, y_v1, z_v1, self.softmax, self.num_classes)
_, _, _, pred_vector2 = utils.classify2vector(
x_v2, y_v2, z_v2, self.softmax, self.num_classes)
_, _, _, pred_vector3 = utils.classify2vector(
x_v3, y_v3, z_v3, self.softmax, self.num_classes)
#visualize vectors
self.draw_vectors(pred_vector1[0].cpu().tolist(),
pred_vector2[0].cpu().tolist(),
pred_vector3[0].cpu().tolist(), draw_img, center,
w)
for i in range(10):
print('Accuracy of %5s : %2d %%' %
(classes[i], 100 * class_correct[i] / class_total[i]))
if __name__ == '__main__':
# Loading and normalizing CIFAR10
trainloader, testloader, classes = data_prepare()
# Define a Convolutional Neural Network
net = ResNet()
net = net.cuda()
x = torch.autograd.Variable(torch.rand(64, 3, 32, 32))
writer.add_graph(net, x.cuda(), verbose=True)
if args.load_model == True:
net.load_state_dict(torch.load(args.model_path))
# Define a Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(),
lr=1e-4,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.,
amsgrad=True)
# Train the network
for epoch in range(args.epochs):
train(net, criterion, optimizer, trainloader, epoch)
test(net, criterion, testloader, epoch)
model_path = './cifar_net.pth'
torch.save(net.state_dict(), model_path)
logger.info(
'========================Validation on valid set====================')
# 读取checkpoint保存的模型，在验证集上跑一遍，计算准确率和召回率
path = './models/ckpoint_' + name + '.tar'
checkpoint = torch.load(path)
# input_shape = (c.hoa_num * 2, c.frames_per_block+2, c.n_freq)
if DATA_TYPE == 'hoa':
block = ResBlock(128, 128, bottleneck=bottle)
else:
block = ResBlock(256, 256, bottleneck=bottle)
model = ResNet(block,
numOfResBlock=num_of_res_blocks,
input_shape=input_shape,
data_type=DATA_TYPE).to(device)
model.load_state_dict(checkpoint['state_dict'])
criterion = nn.MSELoss()
model.eval()
f = open(name + '.txt', 'w')
file_ = open('anechoic_mono_speech_tt.flist', 'r')
all_test_files = file_.readlines()
with torch.no_grad():
for snr in snr_list:
offset = index_offset_dict[snr]
total_correct = np.zeros(numOfEth) # 每个误差容忍度都对应一个准确个数
total_recall = np.zeros(numOfEth)
total_peaks, total_predict = 0, 0 # 总共的真实峰值数，总共预测出来的峰值数
valid_step_cnt = 0
total_valid_loss = 0.0
no_peak = np.zeros(c.scan_num) # 记录无峰值输出（全0输出）的样本个数
x = normalization(x)
x = preprocess_image(x, resize_im=False)
x = recreate_image(x)
if not os.path.exists('./deconv/'+ png_dir):
os.makedirs('./deconv/'+ png_dir)
im_path = './deconv/'+ png_dir+ '/layer_vis_' + str(demode) +'_' + str(index)+ '.jpg'
save_image(x, im_path)
if __name__ == '__main__':
#get model and data
net = ResNet()
net.eval()
net = net.cuda()
net.load_state_dict(torch.load('./cifar_net.pth'))
# print(net.state_dict().keys())
# for name, module in net._modules.items():
# for name, module in module._modules.items():
# print(name)
# for name, param in net.named_parameters():
# print(name)
# net2 = ResNet_deconv(demode=1)
# net2 = net2.cuda()
# encorder = ResNet_encorder(demode=2)
# encorder = encorder.cuda()
# decorder = ResNet_decorder2(demode=2)
# decorder = decorder.cuda()
# params=net.state_dict()
# for k,v in params.items():
# print(k)
if __name__ == '__main__':
args = parse_args()
utils.mkdir(args.save_dir)
# cls and sord
print("Creating model......")
if args.model_name == "mobilenetv2":
model = MobileNetV2(num_classes=args.num_classes)
else:
model = ResNet(torchvision.models.resnet50(pretrained=False),
args.num_classes)
print("Loading weight......")
saved_state_dict = torch.load(args.snapshot)
model.load_state_dict(saved_state_dict)
model.cuda(0)
model.eval() # Change model to 'eval' mode (BN uses moving mean/var).
softmax = nn.Softmax(dim=1).cuda(0)
# test dataLoader
test_loader = loadData(args.test_data, args.input_size, args.batch_size,
args.num_classes, False)
# testing
print('Start testing......')
if args.collect_score:
utils.mkdir(os.path.join(args.save_dir, "collect_score"))
class ResNetPredictor:
def __init__(self, model_path=None):
"""
Params:
model_path: Optional pretrained model file
"""
# Initialize model
self.model = ResNet().cuda()
if model_path is not None:
self.model.load_state_dict(torch.load(model_path))
print('Model read from {}.'.format(model_path))
print('Predictor initialized.')
def fit(self, train_dataset_path, valid_dataset_path, model_dir, **training_args):
"""
train_dataset_path: The path to the training dataset.pkl
valid_dataset_path: The path to the validation dataset.pkl
model_dir: The directory to save models for each epoch
training_args:
- batch_size
- valid_batch_size
- epoch
- lr
- save_every_epoch
"""
# Set paths
self.train_dataset_path = train_dataset_path
self.valid_dataset_path = valid_dataset_path
self.model_dir = model_dir
Path(self.model_dir).mkdir(parents=True, exist_ok=True)
# Set training params
self.batch_size = training_args['batch_size']
self.valid_batch_size = training_args['valid_batch_size']
self.epoch = training_args['epoch']
self.lr = training_args['lr']
self.save_every_epoch = training_args['save_every_epoch']
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
self.onset_criterion = nn.BCEWithLogitsLoss()
self.offset_criterion = nn.BCEWithLogitsLoss()
self.pitch_criterion = nn.CrossEntropyLoss()
# Read the datasets
print('Reading datasets...')
with open(self.train_dataset_path, 'rb') as f:
self.training_dataset = pickle.load(f)
with open(self.valid_dataset_path, 'rb') as f:
self.validation_dataset = pickle.load(f)
# Setup dataloader and initial variables
self.train_loader = DataLoader(
self.training_dataset,
batch_size=self.batch_size,
num_workers=4,
pin_memory=True,
shuffle=True,
drop_last=True,
)
self.valid_loader = DataLoader(
self.validation_dataset,
batch_size=self.valid_batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False,
)
start_time = time.time()
training_loss_list = []
valid_loss_list = []
# Start training
self.iters_per_epoch = len(self.train_loader)
for epoch in range(1, self.epoch + 1):
self.model.train()
# Run iterations
total_training_loss = 0
for batch_idx, batch in enumerate(self.train_loader):
self.optimizer.zero_grad()
# Parse batch data
input_tensor = batch[0].permute(0, 2, 1).unsqueeze(1).cuda()
osnet_prob, offset_prob, pitch_class = batch[1][:, 0].float().cuda(), batch[1][:, 1].float().cuda(), batch[1][:, 2].cuda()
# Forward model
onset_logits, offset_logits, pitch_logits = self.model(input_tensor)
# Calculate loss
loss = self.onset_criterion(onset_logits, osnet_prob) \
+ self.offset_criterion(offset_logits, offset_prob) \
+ self.pitch_criterion(pitch_logits, pitch_class)
loss.backward()
self.optimizer.step()
total_training_loss += loss.item()
# Free GPU memory
# torch.cuda.empty_cache()
if epoch % self.save_every_epoch == 0:
# Perform validation
self.model.eval()
with torch.no_grad():
total_valid_loss = 0
for batch_idx, batch in enumerate(self.valid_loader):
# Parse batch data
input_tensor = batch[0].permute(0, 2, 1).unsqueeze(1).cuda()
osnet_prob, offset_prob, pitch_class = batch[1][:, 0].float().cuda(), batch[1][:, 1].float().cuda(), batch[1][:, 2].cuda()
# Forward model
onset_logits, offset_logits, pitch_logits = self.model(input_tensor)
# Calculate loss
loss = self.onset_criterion(onset_logits, osnet_prob) \
+ self.offset_criterion(offset_logits, offset_prob) \
+ self.pitch_criterion(pitch_logits, pitch_class)
total_valid_loss += loss.item()
# Free GPU memory
# torch.cuda.empty_cache()
# Save model
save_dict = self.model.state_dict()
target_model_path = Path(self.model_dir) / 'e_{}'.format(epoch)
torch.save(save_dict, target_model_path)
# Save loss list
training_loss_list.append((epoch, total_training_loss/self.iters_per_epoch))
valid_loss_list.append((epoch, total_valid_loss/len(self.valid_loader)))
# Epoch statistics
print(
'| Epoch [{:4d}/{:4d}] Train Loss {:.4f} Valid Loss {:.4f} Time {:.1f}'.format(
epoch,
self.epoch,
training_loss_list[-1][1],
valid_loss_list[-1][1],
time.time()-start_time,
)
)
# Save loss to file
with open('./plotting/data/loss.pkl', 'wb') as f:
pickle.dump({'train': training_loss_list, 'valid': valid_loss_list}, f)
print('Training done in {:.1f} minutes.'.format((time.time()-start_time)/60))
def _parse_frame_info(self, frame_info):
"""Parse frame info [(onset_probs, offset_probs, pitch_class)...] into desired label format."""
onset_thres = 0.25
offset_thres = 0.25
result = []
current_onset = None
pitch_counter = Counter()
last_onset = 0.0
for idx, info in enumerate(frame_info):
current_time = FRAME_LENGTH*idx + FRAME_LENGTH/2
if info[0] >= onset_thres: # If is onset
if current_onset is None:
current_onset = current_time
last_onset = info[0]
elif info[0] >= onset_thres:
# If current_onset exists, make this onset a offset and the next current_onset
if pitch_counter.most_common(1)[0][0] != 49:
result.append([current_onset, current_time, pitch_counter.most_common(1)[0][0] + 36])
elif len(pitch_counter.most_common(2)) == 2:
result.append([current_onset, current_time, pitch_counter.most_common(2)[1][0] + 36])
current_onset = current_time
last_onset = info[0]
pitch_counter.clear()
elif info[1] >= offset_thres: # If is offset
if current_onset is not None:
if pitch_counter.most_common(1)[0][0] != 49:
result.append([current_onset, current_time, pitch_counter.most_common(1)[0][0] + 36])
elif len(pitch_counter.most_common(2)) == 2:
result.append([current_onset, current_time, pitch_counter.most_common(2)[1][0] + 36])
current_onset = None
pitch_counter.clear()
# If current_onset exist, add count for the pitch
if current_onset is not None:
pitch_counter[info[2]] += 1
return result
def predict(self, test_dataset):
"""Predict results for a given test dataset."""
# Setup params and dataloader
batch_size = 500
test_loader = DataLoader(
test_dataset,
batch_size=batch_size,
pin_memory=False,
shuffle=False,
drop_last=False,
)
# Start predicting
results = []
self.model.eval()
with torch.no_grad():
print('Forwarding model...')
song_frames_table = {}
for batch_idx, batch in enumerate(tqdm(test_loader)):
# Parse batch data
input_tensor = batch[0].unsqueeze(1).cuda()
song_ids = batch[1]
# Forward model
onset_logits, offset_logits, pitch_logits = self.model(input_tensor)
onset_probs, offset_probs, pitch_logits = torch.sigmoid(onset_logits).cpu(), torch.sigmoid(offset_logits).cpu(), pitch_logits.cpu()
# Collect frames for corresponding songs
for bid, song_id in enumerate(song_ids):
frame_info = (onset_probs[bid], offset_probs[bid], torch.argmax(pitch_logits[bid]).item())
song_frames_table.setdefault(song_id, [])
song_frames_table[song_id].append(frame_info)
# Parse frame info into output format for every song
results = {}
for song_id, frame_info in song_frames_table.items():
results[song_id] = self._parse_frame_info(frame_info)
return results
