Spaces:
Running on Zero
Running on Zero
File size: 2,677 Bytes
d43892c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | from argparse import ArgumentParser
import torch
from models.trainer import *
print(torch.cuda.is_available())
"""
the main function for training the CD networks
"""
def train(args):
dataloaders = utils.get_loaders(args)
model = CDTrainer(args=args, dataloaders=dataloaders)
model.train_models()
def test(args):
from models.evaluator import CDEvaluator
dataloader = utils.get_loader(args.data_name, img_size=args.img_size,
batch_size=args.batch_size, is_train=False,
split='test')
model = CDEvaluator(args=args, dataloader=dataloader)
model.eval_models()
if __name__ == '__main__':
# ------------
# args
# ------------
parser = ArgumentParser()
parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
parser.add_argument('--project_name', default='test', type=str)
parser.add_argument('--checkpoint_root', default='checkpoints', type=str)
# data
parser.add_argument('--num_workers', default=4, type=int)
parser.add_argument('--dataset', default='CDDataset', type=str)
parser.add_argument('--data_name', default='LEVIR', type=str)
parser.add_argument('--batch_size', default=8, type=int)
parser.add_argument('--split', default="train", type=str)
parser.add_argument('--split_val', default="val", type=str)
parser.add_argument('--img_size', default=256, type=int)
# model
parser.add_argument('--n_class', default=2, type=int)
parser.add_argument('--net_G', default='base_transformer_pos_s4_dd8', type=str,
help='base_resnet18 | base_transformer_pos_s4 | '
'base_transformer_pos_s4_dd8 | '
'base_transformer_pos_s4_dd8_dedim8|')
parser.add_argument('--loss', default='ce', type=str)
# optimizer
parser.add_argument('--optimizer', default='sgd', type=str)
parser.add_argument('--lr', default=0.01, type=float)
parser.add_argument('--max_epochs', default=100, type=int)
parser.add_argument('--lr_policy', default='linear', type=str,
help='linear | step')
parser.add_argument('--lr_decay_iters', default=100, type=int)
args = parser.parse_args()
utils.get_device(args)
print(args.gpu_ids)
# checkpoints dir
args.checkpoint_dir = os.path.join(args.checkpoint_root, args.project_name)
os.makedirs(args.checkpoint_dir, exist_ok=True)
# visualize dir
args.vis_dir = os.path.join('vis', args.project_name)
os.makedirs(args.vis_dir, exist_ok=True)
train(args)
test(args)
|