from models import ModelBuilder, activate
from utils import AverageMeter, \
recover_rgb, magnitude2heatmap, \
istft_reconstruction, warpgrid, \
combine_video_audio, save_video, makedirs
from viz import plot_loss_metrics, HTMLVisualizer
import imageio
def main(args):
dataset_train = MUSICMixDataset(args.list_train, args, split='train')
loader_train = torch.utils.data.DataLoader(dataset_train,
batch_size=20,
shuffle=True,
num_workers=0,
drop_last=True)
print("a")
for data in loader_train:
print('new iter')
print(data)
print("b")
if __name__ == '__main__':
parser = ArgParser()
args = parser.parse_train_arguments()
args.batch_size = args.num_gpus * args.batch_size_per_gpu
args.device = torch.device("cuda")
main(args)
from torch import nn
from arguments import ArgParser
from constants import *
from data import *
from logger import *
import models
from training_tools import Evaluator, Optimizer, Predictor
from training_tools import find_save_dir, save_progress, load_progress
from utils import *
# Define cloud storage here
cloudFS = GCStorage.get_CloudFS(PROJECT_ID, GC_BUCKET, CREDENTIAL_PATH)
# Parse command line arguments
parser = ArgParser()
args = parser.parse_args()
# Verify arguments
# check_args(args)
# Setup logger
logger = args.misc_args.logger
# Setup device
device = args.misc_args.device
# Select columns
selected_columns = sorted(select_columns(args.data_args.data_spec))
if args.model_args.wavelet:
args.model_args.input_channels = len(selected_columns) * 2