diff options
| author | SsnL <tongzhou.wang.1994@gmail.com> | 2018-01-13 23:04:43 -0500 |
|---|---|---|
| committer | SsnL <tongzhou.wang.1994@gmail.com> | 2018-01-13 23:04:43 -0500 |
| commit | 1c5c2f50da5ae101077c27fdac2a12fb1619ec86 (patch) | |
| tree | cec1212a4e4327a60a029f1e04ace0a48fb22327 | |
| parent | 929454c133fb19c03264b00179dae25f458efd36 (diff) | |
fix resize_ issue #170
| -rw-r--r-- | models/cycle_gan_model.py | 9 | ||||
| -rw-r--r-- | models/pix2pix_model.py | 12 | ||||
| -rw-r--r-- | models/test_model.py | 6 |
3 files changed, 13 insertions, 14 deletions
diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py index fe06823..b7b840d 100644 --- a/models/cycle_gan_model.py +++ b/models/cycle_gan_model.py @@ -20,8 +20,6 @@ class CycleGANModel(BaseModel): nb = opt.batchSize size = opt.fineSize - self.input_A = self.Tensor(nb, opt.input_nc, size, size) - self.input_B = self.Tensor(nb, opt.output_nc, size, size) # load/define networks # The naming conversion is different from those used in the paper @@ -81,8 +79,11 @@ class CycleGANModel(BaseModel): AtoB = self.opt.which_direction == 'AtoB' input_A = input['A' if AtoB else 'B'] input_B = input['B' if AtoB else 'A'] - self.input_A.resize_(input_A.size()).copy_(input_A) - self.input_B.resize_(input_B.size()).copy_(input_B) + if len(self.gpu_ids) > 0: + input_A = input_A.cuda(self.gpu_ids[0], async=True) + input_B = input_B.cuda(self.gpu_ids[0], async=True) + self.input_A = input_A + self.input_B = input_B self.image_paths = input['A_paths' if AtoB else 'B_paths'] def forward(self): diff --git a/models/pix2pix_model.py b/models/pix2pix_model.py index 56adfc1..74a941e 100644 --- a/models/pix2pix_model.py +++ b/models/pix2pix_model.py @@ -16,11 +16,6 @@ class Pix2PixModel(BaseModel): def initialize(self, opt): BaseModel.initialize(self, opt) self.isTrain = opt.isTrain - # define tensors - self.input_A = self.Tensor(opt.batchSize, opt.input_nc, - opt.fineSize, opt.fineSize) - self.input_B = self.Tensor(opt.batchSize, opt.output_nc, - opt.fineSize, opt.fineSize) # load/define networks self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, @@ -64,8 +59,11 @@ class Pix2PixModel(BaseModel): AtoB = self.opt.which_direction == 'AtoB' input_A = input['A' if AtoB else 'B'] input_B = input['B' if AtoB else 'A'] - self.input_A.resize_(input_A.size()).copy_(input_A) - self.input_B.resize_(input_B.size()).copy_(input_B) + if len(self.gpu_ids) > 0: + input_A = input_A.cuda(self.gpu_ids[0], async=True) + input_B = input_B.cuda(self.gpu_ids[0], async=True) + self.input_A = input_A + self.input_B = input_B self.image_paths = input['A_paths' if AtoB else 'B_paths'] def forward(self): diff --git a/models/test_model.py b/models/test_model.py index 2ae2812..f593c46 100644 --- a/models/test_model.py +++ b/models/test_model.py @@ -12,8 +12,6 @@ class TestModel(BaseModel): def initialize(self, opt): assert(not opt.isTrain) BaseModel.initialize(self, opt) - self.input_A = self.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize) - self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, @@ -29,7 +27,9 @@ class TestModel(BaseModel): def set_input(self, input): # we need to use single_dataset mode input_A = input['A'] - self.input_A.resize_(input_A.size()).copy_(input_A) + if len(self.gpu_ids) > 0: + input_A = input_A.cuda(self.gpu_ids[0], async=True) + self.input_A = input_A self.image_paths = input['A_paths'] def test(self): |