train function
(train_loader, model, optimizer, epoch, save_path, writer)
| 66 | return loss.mean() |
| 67 | |
| 68 | def train(train_loader, model, optimizer, epoch, save_path, writer): |
| 69 | """ |
| 70 | train function |
| 71 | """ |
| 72 | global step |
| 73 | model.train() |
| 74 | loss_all = 0 |
| 75 | epoch_step = 0 |
| 76 | try: |
| 77 | for i, (images, gts, scribble, gt_smalls) in enumerate(train_loader, start=1): |
| 78 | optimizer.zero_grad() |
| 79 | images = images.cuda(device=device_ids[0]) |
| 80 | gts = gts.cuda(device=device_ids[0]) |
| 81 | scribble = scribble.cuda(device=device_ids[0]) |
| 82 | # edges = edges.cuda(device=device_ids[0]) |
| 83 | gt_smalls = gt_smalls.cuda(device = device_ids[0]) |
| 84 | |
| 85 | preds = model(images) |
| 86 | |
| 87 | # ual loss |
| 88 | ual_coef = get_coef(iter_percentage=i/total_step, method='cos') |
| 89 | ual_loss = cal_ual(seg_logits=preds[4], seg_gts=gts) |
| 90 | ual_loss *= ual_coef |
| 91 | |
| 92 | # structure loss w/o weighted |
| 93 | if opt.lossmanner == 'NoWeighted': |
| 94 | loss_init = structure_loss(preds[0], gts)*0.0625 |
| 95 | |
| 96 | loss_body = structure_loss(preds[1], gts)*0.125 + structure_loss(preds[2], gts)*0.25 + \ |
| 97 | structure_loss(preds[3], gts)*0.5 |
| 98 | |
| 99 | loss_final = structure_loss(preds[4], gts) |
| 100 | |
| 101 | if opt.lossmanner == 'Weighted': |
| 102 | # weighted sturcture loss |
| 103 | loss_init = weighted_structure_loss(preds[0], gts)*0.0625 |
| 104 | |
| 105 | loss_body = weighted_structure_loss(preds[1], gts)*0.125 + weighted_structure_loss(preds[2], gts)*0.25 + \ |
| 106 | weighted_structure_loss(preds[3], gts)*0.5 |
| 107 | |
| 108 | loss_final = weighted_structure_loss(preds[4], gts) |
| 109 | |
| 110 | if not (bool(opt.lossmanner == 'NoWeighted') | bool(opt.lossmanner == 'Weighted')): |
| 111 | print("Please provide a valid key words for --lossmanner, i.e., 'NoWeighted' or 'Weighted'.") |
| 112 | |
| 113 | |
| 114 | # PCE loss |
| 115 | loss_init1 = partial_cross_entropy(preds[0], scribble.unsqueeze(1))*0.0625 |
| 116 | loss_body1 = partial_cross_entropy(preds[1], scribble.unsqueeze(1))*0.125 + partial_cross_entropy(preds[2], scribble.unsqueeze(1))*0.25 + \ |
| 117 | partial_cross_entropy(preds[3], scribble.unsqueeze(1))*0.5 |
| 118 | loss_final1 = partial_cross_entropy(preds[4], scribble.unsqueeze(1)) |
| 119 | |
| 120 | loss1 = loss_init + loss_body + loss_final |
| 121 | loss2 = loss_init1 + loss_body1 + loss_final1 |
| 122 | loss3 = 2*ual_loss |
| 123 | loss = loss1 + 3*loss2 + loss3 |
| 124 | |
| 125 | loss.backward() |
no test coverage detected