-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain.py
103 lines (81 loc) · 3.17 KB
/
train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
from __future__ import absolute_import
import torch
import torch.nn as nn
from trainer import train
from model.networks import RNN, VideoDiscriminator, ImageDiscriminator
import torchvision
import transforms_vid
from dataset import WZM
#from torch.utils.tensorboard import SummaryWriter
import os
import cfg
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
def main():
args = cfg.parse_args()
torch.cuda.manual_seed(args.random_seed)
print(args)
# create logging folder
log_path = os.path.join(args.save_path, args.exp_name + '/log')
model_path = os.path.join(args.save_path, args.exp_name + '/models')
if not os.path.exists(log_path) and not os.path.exists(model_path):
os.makedirs(log_path)
os.makedirs(model_path)
#writer = SummaryWriter(log_path) # tensorboard
# load model
device = torch.device("cuda:0")
G = RNN(args).to(device)
VD = VideoDiscriminator(args.ch_d).to(device)
ID = ImageDiscriminator(args.ch_d).to(device)
G = nn.DataParallel(G)
VD = nn.DataParallel(VD)
ID = nn.DataParallel(ID)
# Load pre-trained model
pre_trained_epoch = 0
if args.pre_trained_model_path != '':
print(args.pre_trained_model_path)
pre_trained_epoch = 3500
G.load_state_dict(torch.load(os.path.join(args.pre_trained_model_path, 'G_%d.pth'%(pre_trained_epoch))))
VD.load_state_dict(torch.load(os.path.join(args.pre_trained_model_path, 'VD_%d.pth'%(pre_trained_epoch))))
ID.load_state_dict(torch.load(os.path.join(args.pre_trained_model_path, 'ID_%d.pth'%(pre_trained_epoch))))
# optimizer
optimizer_G = torch.optim.Adam(G.parameters(), args.g_lr, (0.5, 0.999))
optimizer_VD = torch.optim.Adam(VD.parameters(), args.d_lr, (0.5, 0.999))
optimizer_ID = torch.optim.Adam(ID.parameters(), args.d_lr, (0.5, 0.999))
# loss
criterion = nn.BCEWithLogitsLoss().to(device)
# prepare dataset
print('==> preparing dataset')
transform = torchvision.transforms.Compose([
#transforms_vid.ClipResize((args.img_height, args.img_height)),
#transforms_vid.ClipCenterCrop(args.img_size),
transforms_vid.ClipToTensor(),
transforms_vid.ClipNormalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))]
)
dataset = WZM(args.data_path, transform=transform)
dataloader = torch.utils.data.DataLoader(
dataset = dataset,
batch_size = args.batch_size,
num_workers = args.num_workers,
shuffle = True,
pin_memory = True,
drop_last = True
)
# for validation
fixed_zfg = torch.randn(args.n_za_test, args.d_za, 1, 1, 1).to(device)
#fixed_zbg = torch.randn(args.n_za_test, args.d_za, 1, 1, 1).to(device)
#fixed_zm = torch.randn(args.n_zm_test, args.d_zm, 1, 1, 1).to(device)
print('==> start training')
for epoch in range(pre_trained_epoch, args.max_epoch):
train(args, epoch, G, VD, ID, optimizer_G, optimizer_VD, optimizer_ID, criterion, dataloader, device)
'''
if epoch % args.val_freq == 0:
vis(epoch, G, fixed_zfg, fixed_zbg, fixed_zm, writer, device)
'''
if epoch % args.save_freq == 0:
torch.save(G.state_dict(), os.path.join(model_path, 'G_%d.pth'%(epoch)))
torch.save(VD.state_dict(), os.path.join(model_path, 'VD_%d.pth'%(epoch)))
torch.save(ID.state_dict(), os.path.join(model_path, 'ID_%d.pth'%(epoch)))
return
if __name__ == '__main__':
main()