def _calculate_pixelwise_loss(self, hm, target_coords):
sigma = 1.0
if self.pixelwise_loss == 'jsd':
return js_reg_losses(hm, target_coords, sigma)
elif self.pixelwise_loss is None:
return 0
raise Exception('unrecognised pixelwise loss: {}'.format(self.pixelwise_loss))
def train_new_data_with_model():
model = CoordRegression(n_locations=8)
optimizer = optim.RMSprop(model.parameters(), lr=2.5e-4, alpha=0.9)
model = torch.nn.DataParallel(model).cuda()
from data_process_landmarks_hw import dataloader
# 训练集
dataloader = dataloader
for epoch in range(10):
epoch_start = time.time()
print("Epoch: {}/{}".format(epoch + 1, 10))
train_loss = 0.0
train_loss_cord = []
# forward pass
model.train()
# 训练
for i_batch, data in enumerate(dataloader):
img, landmarks = data
img = torch.tensor(img, dtype=torch.float32)
img = img.to(device)
landmarks = torch.tensor(landmarks / 64.0, dtype=torch.float32)
landmarks = landmarks.to(device)
# print("Ground-truth:", gt_hmap.shape)
optimizer.zero_grad()
# forward pass
coords, heatmaps = model(img)
# per-location euclidean losses
euc_losses = dsntnn.euclidean_losses(coords, landmarks)
# print("predict coords", coords, landmarks)
# per-location regulation losses
reg_losses = dsntnn.js_reg_losses(heatmaps, landmarks, sigma_t=1.0)
# combine losses into an overall loss
loss = dsntnn.average_loss(euc_losses + reg_losses)
# calculate gradients
optimizer.zero_grad()
loss.backward()
# update model parameters with RMSprop
optimizer.step()
train_loss_cord.append(loss)
if i_batch % 20 == 19:
print(loss, euc_losses, reg_losses)
# break
# print(loss)
torch.save(
model,
'models/' + 'landmarks' + '_model_new_data_8' + str(epoch) + '.pt')
print(train_loss_cord)
def _calculate_reg_loss(self, target_var, reg, hm_var, hm_sigma):
# Convert sigma (aka standard deviation) from pixels to normalized units
sigma = (2.0 * hm_sigma / hm_var.size(-1))
# Apply a regularisation term relating to the shape of the heatmap.
if reg == 'var':
reg_loss = dsntnn.variance_reg_losses(hm_var, sigma, )
elif reg == 'kl':
reg_loss = dsntnn.kl_reg_losses(hm_var, target_var, sigma)
elif reg == 'js':
reg_loss = dsntnn.js_reg_losses(hm_var, target_var, sigma)
# elif reg == 'mse':
# reg_loss = dsntnn.mse_reg_losses(hm_var, target_var, sigma, mask_var)
else:
reg_loss = 0
return reg_loss
def calc_coord_loss(coords, heatmaps, target_var, masks):
# Per-location euclidean losses
euc_losses = dsntnn.euclidean_losses(
coords,
target_var) # shape:[B, D, L, 2] batch, depth, locations, feature
# Per-location regularization losses
reg_losses = []
for i in range(heatmaps.shape[1]):
hms = heatmaps[:, i]
target = target_var[:, i]
reg_loss = dsntnn.js_reg_losses(hms, target, sigma_t=1.0)
reg_losses.append(reg_loss)
reg_losses = torch.stack(reg_losses, 1)
# reg_losses = dsntnn.js_reg_losses(heatmaps, target_var, sigma_t=1.0) # shape: [B, D, L, 7, 7]
# Combine losses into an overall loss
coord_loss = dsntnn.average_loss((euc_losses + reg_losses).squeeze(),
mask=masks)
return coord_loss
def forward(self, heatmaps, coords, targets, device):
out = torch.Tensor([31, 31]).to(device)
batch_size = coords.shape[0]
n_stages = coords.shape[1]
if len(targets.shape) != len(coords.shape):
targets = torch.unsqueeze(targets, 1)
targets = targets.repeat(1, n_stages, 1, 1)
targets = (targets.div(255) * 2 + 1) / out - 1
losses = []
for i in range(batch_size):
euc_loss = dsntnn.euclidean_losses(coords[i, :, :, :],
targets[i, :, :, :])
reg_loss = dsntnn.js_reg_losses(heatmaps[i, :, :, :, :],
targets[i, :, :, :],
sigma_t=1.0)
losses.append(dsntnn.average_loss(euc_loss + reg_loss))
return sum(losses) / batch_size
def coord_and_heatmap_loss(_coords, _heatmaps, y):
if not isinstance(_coords, list):
coords = [_coords]
heatmaps = [_heatmaps]
else:
coords = _coords
heatmaps = _heatmaps
_losses = []
for c, h in zip(
coords, heatmaps
): # for intermediate supervision: apply loss to all outputs
coords_loss = sum(
[criterion(c, y) for criterion in coord_criterions])
heatmap_losses = dsntnn.js_reg_losses(
h, y, sigma_t=self.config["heatmap_sigma"]
) # TODO different sigmas in each HG?
heatmap_loss = dsntnn.average_loss(heatmap_losses)
_losses.append(
coords_loss + heatmap_loss
) # TODO add lambda for regularization strength??
return sum(_losses)
total_correct = 0
total_acc_misses = np.zeros((len(threshes)), dtype=int)
model.train()
start = time.time()
for images, labels in train_loader: # Get Batch
cuda_images, cuda_labels = images.cuda(main_gpu), labels.cuda(main_gpu)
# Forward pass
coords, heatmaps = model(cuda_images)
total_acc_misses = total_acc_misses + accuracy_func(
coords, cuda_labels)
# Loss
euc_losses = dsntnn.euclidean_losses(coords, cuda_labels)
reg_losses = dsntnn.js_reg_losses(heatmaps, cuda_labels, sigma_t=1.0)
loss = dsntnn.average_loss(euc_losses + reg_losses)
total_loss += loss.item()
# Calculate gradients
optimizer.zero_grad()
loss.backward()
# Update model parameters with RMSprop
optimizer.step()
end = time.time()
torch.save(model.state_dict(),
checkpoints_path + "epoch_{}.pth".format(epoch))
test_correct = 0 #running count of correct predictions on test set
total = 0
for i, data in enumerate(trainloader, 0):
# get the inputs
inputs, labels = data
# load to GPU
inputs = inputs.cuda()
labels = labels.cuda()
# wrap them in Variable
inputs, labels = Variable(inputs), Variable(labels)
#forward pass
coords, heatmaps, xhist = model(inputs)
#Per-location euclidean losses
euc_losses = dsntnn.euclidean_losses(coords, labels)
#Per-location regularization losses
reg_losses = dsntnn.js_reg_losses(heatmaps,labels,sigma_t=1.0)
#Combine losses into an overall loss
loss = dsntnn.average_loss(euc_losses+reg_losses)
#Calculate gradients
optimizer.zero_grad()
torch.nn.utils.clip_grad_norm_(model.parameters(),max_norm=5)
loss.backward()
#Update model parameters
optimizer.step()
# accrue loss
running_train_loss += loss.data[0]
for data in testloader:
testin, testlabels = data
# Load to GPU
testin = testin.cuda()
raccoon_face_tensor = torch.from_numpy(img).permute(2, 0, 1).float()
input_tensor = raccoon_face_tensor.div(255).unsqueeze(0)
input_var = input_tensor.cuda()
eye_coords_tensor = torch.Tensor([[label_all]])
target_tensor = (eye_coords_tensor * 2 +
1) / torch.Tensor(image_size) - 1
target_var = target_tensor.cuda()
coords, heatmaps = model(input_var)
# Per-location euclidean losses
euc_losses = dsntnn.euclidean_losses(coords, target_var)
# Per-location regularization losses
reg_losses = dsntnn.js_reg_losses(heatmaps, target_var,
sigma_t=1.0).cuda()
# Combine losses into an overall loss
loss = dsntnn.average_loss(euc_losses + reg_losses).cuda()
# Calculate gradients
optimizer.zero_grad()
loss.backward()
count += 1
if count % 200 == 0:
print("process: " + str(count) + " /2000 in epoch: " +
str(i) + str(target_var))
print("loss: " + str(loss) + " coords: " +
str(list(coords.data[0, 0])))
logging.info("process: " + str(count) + " /2000 in epoch: " +
str(i) + str(target_var))
# labels = labels.cuda()
# clean=clean.cuda()
# wrap them in Variable
inputs,labels,clean= Variable(inputs),Variable(labels),Variable(clean)
#forward pass
if mode=='single':
coords, heatmaps,_,_= model(inputs.float())
else:
if mode=='WN-less':
coords, heatmaps,clean_outputs,_= model(inputs.float())
else:
coords, heatmaps,clean_outputs,ridge,noise= model(inputs.float())
#Per-location euclidean losses
euc_losses = dsntnn.euclidean_losses(coords, labels)
#Per-location regularization losses
reg_losses = dsntnn.js_reg_losses(heatmaps,labels,sigma_t=1)
if mode=='single':
loss = dsntnn.average_loss(euc_losses+reg_losses)
else:
#fidelity loss
fid_loss=criterion(clean_outputs,clean.float())
if mode=='WN-less':
loss = dsntnn.average_loss(euc_losses+reg_losses)+.5*fid_loss
else:
ridge_tar=torch.zeros([3,4,ridge.shape[2],ridge.shape[3]])
nplabels=torch.zeros([3,4,2])
#extracting ridge_clean
for l in range(3):
for m in range(4):
nplabels[l, m, 0] = labels[l, m, 0] * ((40) / 2 / .9) + (
(40) / 2 + 10)
