| 58 | |
| 59 | |
| 60 | def build_segmenter(args, DDP=True, OPEN=False): |
| 61 | model = build_model(args) |
| 62 | if DDP: |
| 63 | if args.sync_bn: |
| 64 | model = nn.SyncBatchNorm.convert_sync_batchnorm(model) |
| 65 | model = nn.parallel.DistributedDataParallel(model.cuda(), |
| 66 | device_ids=[args.gpu], |
| 67 | find_unused_parameters=True |
| 68 | ) |
| 69 | |
| 70 | single_model = model.module |
| 71 | if OPEN: |
| 72 | for p in single_model.backbone.parameters(): |
| 73 | p.requires_grad_(False) |
| 74 | param_list = [ |
| 75 | { |
| 76 | "params": [ |
| 77 | p |
| 78 | for n, p in single_model.named_parameters() |
| 79 | if "backbone" not in n and "text_encoder" not in n and p.requires_grad |
| 80 | ], |
| 81 | |
| 82 | }, |
| 83 | { |
| 84 | "params": [ |
| 85 | p |
| 86 | for n, p in single_model.named_parameters() |
| 87 | if "pwam" in n and p.requires_grad |
| 88 | ], |
| 89 | |
| 90 | }, |
| 91 | { |
| 92 | "params": [p for n, p in single_model.named_parameters() if "backbone" in n and "pwam" not in n and p.requires_grad], |
| 93 | "lr": args.lr_backbone, |
| 94 | }, |
| 95 | { |
| 96 | "params": [p for n, p in single_model.named_parameters() if "text_encoder" in n and p.requires_grad], |
| 97 | "lr": args.lr_text_encoder, |
| 98 | }, |
| 99 | ] |
| 100 | |
| 101 | return model, param_list |
| 102 | else: |
| 103 | model = nn.DataParallel(model).cuda() |
| 104 | return model |