modify the training process

Author: wbw520

Diff for: configs.py

@@ -6,15 +6,16 @@ def get_args_parser():
 
     # train settings
     parser.add_argument("--dataset", type=str, default="cityscapes")
-    parser.add_argument("--model_name", type=str, default="PSPNet")
+    parser.add_argument("--model_name", type=str, default="Segmenter")
+    parser.add_argument("--pre_model", type=str, default="mae_pre_epoch99_crop640_patch8_ed768_depth12_head12.pt")
     parser.add_argument("--batch_size", type=int, default=1,
                         help="Number of images sent to the network in one step.")
     parser.add_argument("--root", type=str, default="/home/wangbowen/DATA/cityscapes",
                         help="Path to the directory containing the image list.")
     parser.add_argument("--crop_size", type=int, default=[640, 640],
                         help="crop size for training and inference slice.")
     parser.add_argument("--stride_rate", type=float, default=0.5, help="stride ratio.")
-    parser.add_argument("--num_epoch", type=int, default=60, help="Number of training steps.")
+    parser.add_argument("--num_epoch", type=int, default=200, help="Number of training steps.")
     parser.add_argument("--num_classes", type=int, default=19, help="Number of class for dataset.")
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')

Diff for: mae_pre-training.py

@@ -7,7 +7,7 @@
 from torch.utils.data import DistributedSampler
 from data import cityscapes
 from data.loader_tools import get_joint_transformations, get_standard_transformations
-from model import mae_model as models_mae
+from model.mae_model import mae_vit
 from utils2.misc import NativeScalerWithGradNormCount as NativeScaler
 import timm.optim.optim_factory as optim_factory
 import os
@@ -23,17 +23,14 @@
 
 def get_args_parser():
     parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
-    parser.add_argument('--batch_size', default=1, type=int,
+    parser.add_argument('--batch_size', default=4, type=int,
                         help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=200, type=int)
+    parser.add_argument('--num_epochs', default=100, type=int)
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
     parser.add_argument('--use_pre', default=False, type=bool)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch8', type=str, metavar='MODEL',
-                        help='Name of model to train')
-    parser.add_argument('--patch_size', default=8, type=int, help='size of patch')
     parser.add_argument("--crop_size", type=int, default=[640, 640],
                         help="crop size for training and inference slice.")
     parser.add_argument("--stride_rate", type=float, default=0.5, help="stride ratio.")
@@ -45,6 +42,15 @@ def get_args_parser():
                         help='Use (per-patch) normalized pixels as targets for computing loss')
     parser.set_defaults(norm_pix_loss=False)
 
+    # VIT settings
+    parser.add_argument("--patch_size", type=int, default=16, help="define the patch size.")
+    parser.add_argument("--encoder_embed_dim", type=int, default=1024, help="dimension for encoder.")
+    parser.add_argument("--decoder_embed_dim", type=int, default=512, help="dimension for decoder.")
+    parser.add_argument("--encoder_depth", type=int, default=24, help="depth for encoder.")
+    parser.add_argument("--decoder_depth", type=int, default=8, help="depth for decoder.")
+    parser.add_argument("--encoder_num_head", type=int, default=16, help="head number for encoder.")
+    parser.add_argument("--decoder_num_head", type=int, default=16, help="head number for decoder.")
+
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
@@ -56,13 +62,13 @@ def get_args_parser():
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
     parser.add_argument('--root', default="/home/wangbowen/DATA/cityscapes", type=str,
                         help='dataset path')
-    parser.add_argument('--output_dir', default='save_model',
+    parser.add_argument('--output_dir', default='save_model/',
                         help='path where to save, empty for no saving')
     parser.add_argument('--log_dir', default='save_model',
                         help='path where to tensorboard log')
@@ -86,7 +92,7 @@ def main():
     misc.init_distributed_mode(args)
     device = torch.device(args.device)
     cudnn.benchmark = True
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss, img_size=args.crop_size[0])
+    model = mae_vit(args)
     model.to(device)
     model_without_ddp = model
     # print("Model = %s" % str(model_without_ddp))
@@ -97,6 +103,7 @@ def main():
                                       standard_transform=standard_transformations)
 
     if args.use_pre:
+        # use the pre-trained parameter from mae paper
         checkpoint = torch.load("save_model/mae_visualize_vit_large.pth", map_location='cpu')
         checkpoint_model = checkpoint['model']
         interpolate_pos_embed(model, checkpoint_model)
@@ -138,17 +145,18 @@ def main():
     print(optimizer)
     loss_scaler = NativeScaler()
 
-    print(f"Start training for {args.epochs} epochs")
+    print(f"Start training for {args.num_epochs} epochs")
     start_time = time.time()
 
-    for epoch in range(args.epochs):
+    for epoch in range(args.num_epochs):
         if args.distributed:
             sampler_train.set_epoch(epoch)
 
-        if args.output_dir and ((epoch + 1) % 20 == 0 or epoch + 1 == args.epochs):
-            misc.save_model(
-                args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+        if args.output_dir and ((epoch + 1) % 50 == 0 or epoch + 1 == args.num_epochs):
+            torch.save(model_without_ddp.state_dict(),
+                       args.output_dir + "mae_pre_epoch" + str(epoch) + "_crop" + str(args.crop_size[0]) + "_patch" +
+                       str(args.patch_size) + "_ed" + str(args.encoder_embed_dim) + "_depth" + str(args.encoder_depth) +
+                       "_head" + str(args.encoder_num_head) + ".pt")
 
         train_stats = train_one_epoch(
             model, train_loader,

Diff for: model/mae_model.py

@@ -216,13 +216,9 @@ def forward(self, imgs, mask_ratio=0.75):
         return loss, pred, mask
 
 
-def mae_vit(**kwargs):
+def mae_vit(args):
     model = MaskedAutoencoderViT(
-        patch_size=16, embed_dim=768, depth=12, num_heads=12,
-        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
-        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
-    return model
-
-
-# set recommended archs
-mae_vit = mae_vit
+        img_size=args.crop_size[0], patch_size=args.patch_size, embed_dim=args.encoder_embed_dim, depth=args.encoder_depth, num_heads=args.encoder_num_head,
+        decoder_embed_dim=args.decoder_embed_dim, decoder_depth=args.decoder_depth, decoder_num_heads=args.decoder_num_head,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6))
+    return model

Diff for: model/segmenter.py

@@ -1,6 +1,8 @@
 from model.vit_model import vit_encoder
 from model.segmenter_decoder import sg_vit_mask_decoder
 import torch.nn as nn
+import torch
+import os
 import torch.nn.functional as F
 
 
@@ -42,6 +44,17 @@ def get_attention_map_dec(self, im, layer_id):
 
 def create_segmenter(args):
     encoder = vit_encoder(img_size=args.crop_size[0], patch_size=args.patch_size, embed_dim=args.encoder_embed_dim, depth=args.encoder_depth, num_heads=args.encoder_num_head)
+    check_point = torch.load(os.path.join("save_model", args.pre_model), map_location="cuda:0")
+    state_dict = ["decoder", "mask_token"]
+    record = []
+    for k, v in check_point.items():
+        if state_dict[0] in k or state_dict[1] in k:
+            record.append(k)
+    for item in record:
+        del check_point[item]
+    encoder.load_state_dict(check_point, strict=True)
+    print("load pre-trained weight from: ", args.pre_model)
+
     decoder = sg_vit_mask_decoder(patch_size=args.patch_size, encoder_embed_dim=args.encoder_embed_dim,
                                   decoder_embed_dim=args.decoder_embed_dim, decoder_depth=args.decoder_depth, decoder_num_heads=args.decoder_num_head, n_cls=args.num_classes)
     model = Segmenter(encoder, decoder)

Diff for: model/segmenter_decoder.py

@@ -62,10 +62,7 @@ def forward(self, x, im_size):
 
         patches = patches / patches.norm(dim=-1, keepdim=True)
         cls_seg_feat = cls_seg_feat / cls_seg_feat.norm(dim=-1, keepdim=True)
-        print(patches.shape)
-        print(cls_seg_feat.shape)
         masks = patches @ cls_seg_feat.transpose(1, 2)
-        print(masks.shape)
         masks = self.mask_norm(masks)
         masks = rearrange(masks, "b (h w) n -> b n h w", h=int(GS))
 

Diff for: model/vit_model.py

@@ -79,5 +79,5 @@ def resize_pos_embed(posemb, grid_old_shape, grid_new_shape, num_extra_tokens):
 
 
 def vit_encoder(**kwargs):
-    model = VisionTransformer(mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model = VisionTransformer(mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), num_classes=0, **kwargs)
     return model

Diff for: sample_demo/color_mask_translate.py

@@ -7,7 +7,7 @@
 
 
 def main():
-    root = "/home/wangbowen/DATA/facades/facade_raw"
+    root = "/home/wangbowen/DATA/Facade/"
     item_list = get_name(root, mode_folder=False)
     image_list = []
     for item in item_list:

Diff for: train.py

@@ -49,6 +49,7 @@ def main():
     # param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
     param_groups = [p for p in model_without_ddp.parameters() if p.requires_grad]
     optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+
     print(optimizer)
     loss_scaler = NativeScaler()
 

Diff for: utils/engine.py

@@ -10,11 +10,14 @@
 def train_model(args, epoch, model, train_loader, criterion, optimizer, loss_scaler, device):
     model.train()
     train_main_loss = AverageMeter('Train Main Loss', ':.5')
-    train_aux_loss = AverageMeter('Train Aux Loss', ':.5')
     lr = AverageMeter('lr', ':.5')
     L = len(train_loader)
     curr_iter = epoch * L
-    progress = ProgressMeter(L, [train_main_loss, train_aux_loss, lr], prefix="Epoch: [{}]".format(epoch))
+    record = [lr, train_main_loss]
+    if args.model_name == "PSPNet":
+        train_aux_loss = AverageMeter('Train Aux Loss', ':.5')
+        record.append(train_aux_loss)
+    progress = ProgressMeter(L, record, prefix="Epoch: [{}]".format(epoch))
     accum_iter = args.accum_iter
 
     for data_iter_step, data in enumerate(train_loader):
@@ -26,10 +29,15 @@ def train_model(args, epoch, model, train_loader, criterion, optimizer, loss_sca
         mask = data["masks"].to(device, dtype=torch.int64)
 
         with torch.cuda.amp.autocast():
-            outputs, aux = model(inputs)
-            main_loss = criterion(outputs, mask)
-            aux_loss = criterion(aux, mask)
-            loss = main_loss + 0.4 * aux_loss
+            if args.model_name == "PSPNet":
+                outputs, aux = model(inputs)
+                main_loss = criterion(outputs, mask)
+                aux_loss = criterion(aux, mask)
+                loss = main_loss + 0.4 * aux_loss
+            else:
+                outputs = model(inputs)
+                main_loss = criterion(outputs, mask)
+                loss = main_loss
 
         loss_value = loss.item()
         if not math.isfinite(loss_value):
@@ -43,7 +51,8 @@ def train_model(args, epoch, model, train_loader, criterion, optimizer, loss_sca
         torch.cuda.synchronize()
 
         train_main_loss.update(main_loss.item())
-        train_aux_loss.update(aux_loss.item())
+        if args.model_name == "PSPNet":
+            train_aux_loss.update(aux_loss.item())
         lr.update(optimizer.param_groups[0]['lr'])
 
         curr_iter += 1
@@ -77,7 +86,7 @@ def evaluation(args, best_record, epoch, model, model_without_ddp, val_loader, c
         best_record['acc_cls'] = acc_cls
         best_record['mean_iou'] = mean_iou
         if args.output_dir:
-            torch.save(model_without_ddp.state_dict(), args.output_dir + str(epoch) + "_epoch_" + args.model_name + ".pt")
+            torch.save(model_without_ddp.state_dict(), args.output_dir + args.model_name + ".pt")
 
     print('-----------------------------------------------------------------------------------------------------------')
     print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iou %.5f]' % (

Diff for: utils2/lr_sched.py

@@ -13,7 +13,7 @@ def adjust_learning_rate(optimizer, epoch, args):
         lr = args.lr * epoch / args.warmup_epochs 
     else:
         lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
-            (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.num_epoch - args.warmup_epochs)))
+            (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.num_epochs - args.warmup_epochs)))
     for param_group in optimizer.param_groups:
         if "lr_scale" in param_group:
             param_group["lr"] = lr * param_group["lr_scale"]