| import os |
| import re |
| from importlib import import_module |
|
|
| import torch |
| import torch.nn as nn |
| from torch.nn.parallel import DataParallel, DistributedDataParallel |
|
|
| import torch.distributed as dist |
| from torch.nn.utils import parameters_to_vector, vector_to_parameters |
|
|
| from .discriminator import Discriminator |
|
|
| from utils import interact |
|
|
| class Model(nn.Module): |
| def __init__(self, args): |
| super(Model, self).__init__() |
|
|
| self.args = args |
| self.device = args.device |
| self.n_GPUs = args.n_GPUs |
| self.save_dir = os.path.join(args.save_dir, 'models') |
| os.makedirs(self.save_dir, exist_ok=True) |
|
|
| module = import_module('model.' + args.model) |
|
|
| self.model = nn.ModuleDict() |
| self.model.G = module.build_model(args) |
| if self.args.loss.lower().find('adv') >= 0: |
| self.model.D = Discriminator(self.args) |
| else: |
| self.model.D = None |
|
|
| self.to(args.device, dtype=args.dtype, non_blocking=True) |
| self.load(args.load_epoch, path=args.pretrained) |
|
|
| def parallelize(self): |
| if self.args.device_type == 'cuda': |
| if self.args.distributed: |
| Parallel = DistributedDataParallel |
| parallel_args = { |
| "device_ids": [self.args.rank], |
| "output_device": self.args.rank, |
| } |
| else: |
| Parallel = DataParallel |
| parallel_args = { |
| 'device_ids': list(range(self.n_GPUs)), |
| 'output_device': self.args.rank |
| } |
|
|
| for model_key in self.model: |
| if self.model[model_key] is not None: |
| self.model[model_key] = Parallel(self.model[model_key], **parallel_args) |
|
|
| def forward(self, input): |
| return self.model.G(input) |
|
|
| def _save_path(self, epoch): |
| model_path = os.path.join(self.save_dir, 'model-{:d}.pt'.format(epoch)) |
| return model_path |
|
|
| def state_dict(self): |
| state_dict = {} |
| for model_key in self.model: |
| if self.model[model_key] is not None: |
| parallelized = isinstance(self.model[model_key], (DataParallel, DistributedDataParallel)) |
| if parallelized: |
| state_dict[model_key] = self.model[model_key].module.state_dict() |
| else: |
| state_dict[model_key] = self.model[model_key].state_dict() |
|
|
| return state_dict |
|
|
| def load_state_dict(self, state_dict, strict=True): |
| for model_key in self.model: |
| parallelized = isinstance(self.model[model_key], (DataParallel, DistributedDataParallel)) |
| if model_key in state_dict: |
| if parallelized: |
| self.model[model_key].module.load_state_dict(state_dict[model_key], strict) |
| else: |
| self.model[model_key].load_state_dict(state_dict[model_key], strict) |
|
|
| def save(self, epoch): |
| torch.save(self.state_dict(), self._save_path(epoch)) |
|
|
| def load(self, epoch=None, path=None): |
| if path: |
| model_name = path |
| elif isinstance(epoch, int): |
| if epoch < 0: |
| epoch = self.get_last_epoch() |
| if epoch == 0: |
| |
| |
| |
|
|
| return |
|
|
| model_name = self._save_path(epoch) |
| else: |
| raise Exception('no epoch number or model path specified!') |
|
|
| print('Loading model from {}'.format(model_name)) |
| state_dict = torch.load(model_name, map_location=self.args.device) |
| self.load_state_dict(state_dict) |
|
|
| return |
|
|
| def synchronize(self): |
| if self.args.distributed: |
| |
| vector = parameters_to_vector(self.parameters()) |
|
|
| dist.broadcast(vector, 0) |
| if self.args.rank != 0: |
| vector_to_parameters(vector, self.parameters()) |
|
|
| del vector |
|
|
| return |
|
|
| def get_last_epoch(self): |
| model_list = sorted(os.listdir(self.save_dir)) |
| if len(model_list) == 0: |
| epoch = 0 |
| else: |
| epoch = int(re.findall('\\d+', model_list[-1])[0]) |
|
|
| return epoch |
|
|
| def print(self): |
| print(self.model) |
|
|
| return |
|
|
