Commit c29bc1a2 authored by Gabriele Franch's avatar Gabriele Franch
Browse files

auto find latest iteration when pretrained_model is a dir

parent f01e1734
......@@ -4,14 +4,13 @@ import argparse
import numpy as np
from tqdm import tqdm
import os
import glob
import core.trainer as trainer
from cikm_inter_dst_predrnn_run_taasss_utils import get_batcher, padding_taasss, change_taasss_dims
from core.models.model_factory import Model
from core.utils import preprocess
from matplotlib import pyplot as plt
import scipy.misc
from PIL import Image
parser = argparse.ArgumentParser(
description="PyTorch video prediction model - DST PredRNN"
......@@ -132,8 +131,16 @@ def schedule_sampling(eta: float, itr: int):
def wrapper_train(model: Model):
itr = 1
eta = args.sampling_start_value
if args.pretrained_model:
if os.path.isdir(args.pretrained_model):
maxiter = 0
for file in os.listdir(args.pretrained_model):
if file.endswith(".pth"):
maxiter = max([int(file[:-4]), maxiter])
args.pretrained_model = f'{args.pretrained_model}/{maxiter}.pth'
itr, eta = model.load(args.pretrained_model)
iterator = get_batcher(args)
progress_bar = tqdm(range(itr, args.max_iterations + 1))
for itr in progress_bar:
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment