class C3D_TCN(nn.Module):
def __init__(self, tcn_out_channel=64, c3d_path='', tcn_path=''):
super(C3D_TCN, self).__init__()
self.c3d = C3D(in_channels=3)
self.tcn = TCN(245760, [128,128,64,tcn_out_channel]) # 245760 == 128, 983040 == 256, 384000 == 160
self.load_models(c3d_path, tcn_path)
def load_models(self, c3d_path, tcn_path):
if os.path.exists(c3d_path):
self.c3d.load_state_dict(torch.load(c3d_path))
if os.path.exists(tcn_path):
self.tcn.load_state_dict(torch.load(tcn_path))
def save_models(self, c3d_path, tcn_path):
torch.save(self.c3d.state_dict(), c3d_path)
torch.save(self.tcn.state_dict(), tcn_path)
def forward(self, X):
N, WC, RGB, F, W, H = X.shape
shape = [N*WC, RGB, F, W, H]
X = self.c3d(X.reshape(shape))
shape = [N, WC, -1]
X = X.reshape(shape)
X = torch.transpose(X, 1, 2)
X = self.tcn(X)
X = torch.transpose(X, 1, 2)
shape = [N, -1]
X = X.reshape(shape)
return X
# data for generator
high_res = np.load('middle_res_shifted.npy')-0.5
low_res = np.array([np.average(high_res[i].reshape(-1,32),axis=1) for i in range(len(high_res))])
train_gen = low_res[:nsample]
test_gen_low = low_res[nsample:]
test_gen_high = high_res[nsample:]
# construct model
model = TCN(1,chans,kernel_size=5,dropout=0) # 1 input channel
generator = srgan.Generator(10,5) # 32 times upsampling
# load/train model
try:
model.load_state_dict(torch.load(model_file))
print('Load Model')
except:
print('No Model found')
train_tcn = True
try:
generator.load_state_dict(torch.load('generator.pt'))
print("Load Generator Model")
except:
print("Generator Model Not Found")
if train_tcn:
model = model.cuda()