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
for k in range(itrain.size()[2]):
ioutput = model(itrain).cuda() #[batch_size,3*ncomponent,length]
loss = LogMixGaussian(ioutput,itrain,batch_size,n_components,index=k)
loss.backward()
optimizer_tcn.step()
else:
ioutput = model(itrain).cuda() #[batch_size,3*ncomponent,length]
loss = LogMixGaussian(ioutput,itrain,batch_size,n_components)
loss.backward()
optimizer_tcn.step()
totloss = LogMixGaussian(ioutput,itrain,batch_size,n_components)
avg_loss += totloss.data
avg_loss/=(nsample//batch_size)
scheduler_tcn.step()
print(j,avg_loss,scheduler_tcn.get_lr())
torch.save(model.state_dict(), model_file)
if train_generator:
generator.cuda()
model.cuda()
optimizer_gen = optim.Adam(generator.parameters(),lr=lr_gen)
scheduler_gen = optim.lr_scheduler.StepLR(optimizer_gen,step_size=step_size_gen,gamma=gamma_gen)
cov = np.cov(train_gen.T)
invcov = np.linalg.inv(cov)
invcov = torch.from_numpy(invcov).float().cuda()
train_gen = torch.from_numpy(train_gen).float().cuda()
train_gen = train_gen.view(train_gen.size()[0]//batch_size,batch_size,1,train_gen.size()[1])
# Save Model
if val_accuracy > max_accuracy:
# Prepare folder
folder_for_this_accuracy = os.path.join(output_folder, str(val_accuracy))
max_accuracy = val_accuracy
print("Models Saved with accuracy={}".format(max_accuracy))
else:
folder_for_this_accuracy = os.path.join(output_folder, "Latest_{}".format(val_accuracy))
if not os.path.exists(folder_for_this_accuracy):
os.mkdir(folder_for_this_accuracy)
# Save networks
torch.save(c3d.state_dict(), os.path.join(folder_for_this_accuracy, "c3d.pkl"))
torch.save(rn.state_dict(), os.path.join(folder_for_this_accuracy, "rn.pkl"))
torch.save(tcn.state_dict(), os.path.join(folder_for_this_accuracy, "tcn.pkl"))
torch.save(ap.state_dict(), os.path.join(folder_for_this_accuracy, "ap.pkl"))
torch.save(c3d_optim.state_dict(), os.path.join(folder_for_this_accuracy, "c3d_optim.pkl"))
torch.save(rn_optim.state_dict(), os.path.join(folder_for_this_accuracy, "rn_optim.pkl"))
torch.save(tcn_optim.state_dict(), os.path.join(folder_for_this_accuracy, "tcn_optim.pkl"))
torch.save(ap_optim.state_dict(), os.path.join(folder_for_this_accuracy, "ap_optim.pkl"))
torch.save(c3d_scheduler.state_dict(), os.path.join(folder_for_this_accuracy, "c3d_scheduler.pkl"))
torch.save(rn_scheduler.state_dict(), os.path.join(folder_for_this_accuracy, "rn_scheduler.pkl"))
torch.save(tcn_scheduler.state_dict(), os.path.join(folder_for_this_accuracy, "tcn_scheduler.pkl"))
torch.save(ap_scheduler.state_dict(), os.path.join(folder_for_this_accuracy, "ap_scheduler.pkl"))
print("Training Done")
print("Final Accuracy = {}".format(max_accuracy))