Wrap up models with DDP Args: model: pretrained encoder, should be frozen classifier: linear classifier
(self, model, classifier)
| 35 | self.logger.log_value('learning_rate', lr, epoch) |
| 36 | |
| 37 | def wrap_up(self, model, classifier): |
| 38 | """Wrap up models with DDP |
| 39 | |
| 40 | Args: |
| 41 | model: pretrained encoder, should be frozen |
| 42 | classifier: linear classifier |
| 43 | """ |
| 44 | args = self.args |
| 45 | model = model.cuda() |
| 46 | classifier = classifier.cuda() |
| 47 | model.eval() |
| 48 | model = DDP(model, device_ids=[args.gpu]) |
| 49 | classifier = DDP(classifier, device_ids=[args.gpu]) |
| 50 | |
| 51 | return model, classifier |
| 52 | |
| 53 | def load_encoder_weights(self, model): |
| 54 | """load pre-trained weights for encoder |