def __init__(self, opt=None, load_deafult=False):
# PREPARATIONS
if opt:
self.opt = opt
else:
self.opt = Opt()
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
print_options(self.opt)
os.makedirs(self.opt.outf, exist_ok=True)
cudnn.benchmark = True
self.load_model()
if load_deafult:
self.load_data(self.opt.style, self.opt.content)
help='which features to transfer, either r31 or r41')
parser.add_argument('--log_dir',
default='./logs/trainv2',
help='Directory to save the log')
################# PREPARATIONS #################
opt = parser.parse_args()
opt.content_layers = opt.content_layers.split(',')
opt.style_layers = opt.style_layers.split(',')
opt.cuda = torch.cuda.is_available()
if (opt.cuda):
torch.cuda.set_device(opt.gpu_id)
os.makedirs(opt.outf, exist_ok=True)
cudnn.benchmark = True
print_options(opt)
log_dir = Path(opt.log_dir)
log_dir.mkdir(exist_ok=True, parents=True)
writer = SummaryWriter(log_dir=str(log_dir))
################# DATA #################
content_dataset = Dataset(opt.contentPath, opt.loadSize, opt.fineSize)
content_loader_ = torch.utils.data.DataLoader(dataset=content_dataset,
batch_size=opt.batchSize,
shuffle=True,
num_workers=1,
drop_last=True)
content_loader = iter(content_loader_)
style_dataset = Dataset(opt.stylePath, opt.loadSize, opt.fineSize)
style_loader_ = torch.utils.data.DataLoader(dataset=style_dataset,
batch_size=opt.batchSize,
help='path to transferred images')
parser.add_argument("--layer", default="r41",
help='which features to transfer, either r31 or r41')
parser.add_argument('--patch_size', type=int, default=1000, help='patch size')
parser.add_argument('--thumb_size', type=int, default=1024, help='thumbnail size')
parser.add_argument('--style_size', type=int, default=1024, help='style size')
parser.add_argument('--padding', type=int, default=32, help='padding')
parser.add_argument('--test_speed', action="store_true", help='test the speed')
parser.add_argument('--URST', action="store_true", help='use URST framework')
parser.add_argument("--device", type=str, default="cuda", help="device")
parser.add_argument('--resize', type=int, default=0, help='resize')
################# PREPARATIONS #################
args = parser.parse_args()
args.cuda = torch.cuda.is_available()
print_options(args)
os.makedirs(args.outf, exist_ok=True)
content_name = args.content.split("/")[-1].split(".")[0]
style_name = args.style.split("/")[-1].split(".")[0]
device = torch.device(args.device)
################# MODEL #################
if(args.layer == 'r31'):
vgg = encoder3().to(device)
dec = decoder3().to(device)
elif(args.layer == 'r41'):
vgg = encoder4().to(device)
dec = decoder4().to(device)
matrix = MulLayer(args.layer).to(device)
vgg.load_state_dict(torch.load(args.vgg_dir))
dec.load_state_dict(torch.load(args.decoder_dir))