main_signal.py

import numpy as np
import os
import random
import torch
import argparse
from Core.Train import train, test


def seed_everything(seed=6767):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def main():

    seed_everything(seed=6767)
    parser = argparse.ArgumentParser(description='signal completion diffusion model')
    parser.add_argument('--state', type=str, default='train', help='train, test')
    parser.add_argument('--epoch', type=int, default=200, help='training epoch')
    parser.add_argument('--dim', type=int, default=96, help='unet_dim')
    parser.add_argument('--block', nargs='+', type=int, default=[1, 2, 4], help='unet_block')
    parser.add_argument('--dataset', type=str, default='EEG', help='ECG, EMG, EEG')
    parser.add_argument('--batch_size', type=int, default=128, help='batch size')
    parser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    parser.add_argument('--beta_1', type=float, default=1e-3, help='beta_1, 1e-3, 1e-4')
    parser.add_argument('--beta_T', type=float, default=1e-3, help='beta_T, 1e-3, 0.02')
    parser.add_argument('--grad_clip', type=float, default=1., help='grad_clip')
    parser.add_argument('--multiplier', type=float, default=2., help='multiplier')
    parser.add_argument('--missing_type_test', type=str, default='block', help='block, scatter, end')
    parser.add_argument('--missing_rate_test', type=float, default=0.5, help='0.5, 0.3, 0.1')
    parser.add_argument('--T', type=int, default=1000, help='diffusion total steps')
    parser.add_argument('--step_T', type=int, default=5, help='diffusion total steps')
    parser.add_argument('--dropout', type=float, default=0.15, help='dropout')
    parser.add_argument('--save_dir', type=str, default='./Checkpoints/', help='address')
    parser.add_argument('--save_weight_dir', type=str, default='./Checkpoints/', help='address')
    parser.add_argument('--save_gen', type=bool, default=True, help='address')
    parser.add_argument('--device', type=str, default='cuda:0', help='device')
    args = parser.parse_args()

    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    if args.state == "train":
        train(args)
    elif args.state == "test":
        test(args)
    else:
        raise NotImplementedError


if __name__ == '__main__':
    main()