CalvinBackup/neuralchen-SimSwap
1
0
Fork 0
Code Issues 288 Pull Requests 16 Actions Packages Projects Releases 2 Wiki Activity

update

Browse Source
This commit is contained in:
chenxuanhong
2022-04-21 22:39:55 +08:00
parent 7ed12d218f
commit 9f3daca179
3 changed files with 6 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -63,6 +63,7 @@ If you find this project useful, please star it. It is the greatest appreciation
Download the dataset from [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ).
The training script is slightly different from the original version, e.g., we replace the patch discriminator with the projected discriminator, which saves a lot of hardware overhead and achieves slightly better results.
In order to ensure normal training, the batch size must be greater than 1.
- Train 256 models
```
models/base_model.py
-7
View File
@@ -51,13 +51,6 @@ class BaseModel(torch.nn.Module):
save_path = os.path.join(self.save_dir, save_filename)
torch.save(network.state_dict(), save_path)
# helper saving function that can be used by subclasses
def save_network(self, network, network_label, epoch_label, gpu_ids=None):
save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
save_path = os.path.join(self.save_dir, save_filename)
torch.save(network.cpu().state_dict(), save_path)
if len(gpu_ids) and torch.cuda.is_available():
network.cuda()
# helper loading function that can be used by subclasses
def load_network(self, network, network_label, epoch_label, save_dir=''):
train.py
+5 -5
View File
@@ -5,7 +5,7 @@
# Created Date: Monday December 27th 2021
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Thursday, 21st April 2022 8:10:05 pm
# Last Modified: Thursday, 21st April 2022 10:36:48 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University
#############################################################
@@ -44,7 +44,7 @@ class TrainOptions:
self.parser.add_argument('--isTrain', type=str2bool, default='True')
# input/output sizes
self.parser.add_argument('--batchSize', type=int, default=16, help='input batch size')
self.parser.add_argument('--batchSize', type=int, default=2, help='input batch size')
# for displays
self.parser.add_argument('--tag', type=str, default='simswap')
@@ -69,9 +69,9 @@ class TrainOptions:
self.parser.add_argument("--Arc_path", type=str, default='arcface_model/arcface_checkpoint.tar', help="run ONNX model via TRT")
self.parser.add_argument("--total_step", type=int, default=1000000, help='total training step')
self.parser.add_argument("--log_frep", type=int, default=250, help='frequence for printing log information')
self.parser.add_argument("--sample_freq", type=int, default=1000, help='frequence for sampling')
self.parser.add_argument("--model_freq", type=int, default=10000, help='frequence for saving the model')
self.parser.add_argument("--log_frep", type=int, default=10, help='frequence for printing log information')
self.parser.add_argument("--sample_freq", type=int, default=30, help='frequence for sampling')
self.parser.add_argument("--model_freq", type=int, default=40, help='frequence for saving the model')