sithu31296 / semantic-segmentation

SOTA Semantic Segmentation Models in PyTorch
MIT License
830 stars 149 forks source link

KeyError: 'model' when I start training #59

Open Bananaspirit opened 11 months ago

Bananaspirit commented 11 months ago

I want to train DRNET on the cityscapes dataset, here are the configurations for training:

DEVICE          : cpu              # device used for training and evaluation (cpu, cuda, cuda0, cuda1, ...)
SAVE_DIR        : '/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/DDR-NET_output' # output folder name used for saving the model, logs and inference results

MODEL:                                    
  NAME          : DDRNet                                           # name of the model you are using
  BACKBONE      : DDRNet-23slim                                                 # model variant
  PRETRAINED    : '/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/pretrained_models/DDRNet-23slim.pth' # backbone model's weight 

DATASET:
  NAME          : CityScapes                                          # dataset name to be trained with (camvid, cityscapes, ade20k)
  ROOT          : '/media/banana/Linux_disk/Docs/data/cityscapes'                                   # dataset root path
  IGNORE_LABEL  : 255

TRAIN:
  IMAGE_SIZE    : [1024, 1024]    # training image size in (h, w)
  BATCH_SIZE    : 8               # batch size used to train
  EPOCHS        : 1             # number of epochs to train
  EVAL_INTERVAL : 20              # evaluation interval during training
  AMP           : false           # use AMP in training
  DDP           : false           # use DDP training

LOSS:
  NAME          : OhemCrossEntropy          # loss function name (ohemce, ce, dice)
  CLS_WEIGHTS   : false            # use class weights in loss calculation

OPTIMIZER:
  NAME          : adamw           # optimizer name
  LR            : 0.001           # initial learning rate used in optimizer
  WEIGHT_DECAY  : 0.01            # decay rate used in optimizer 

SCHEDULER:
  NAME          : warmuppolylr    # scheduler name
  POWER         : 0.9             # scheduler power
  WARMUP        : 10              # warmup epochs used in scheduler
  WARMUP_RATIO  : 0.1             # warmup ratio

EVAL:
  MODEL_PATH    : '/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/DDR-NET_output/ddrnet_23slim_city.pth'     # trained model file path
  IMAGE_SIZE    : [1024, 1024]                            # evaluation image size in (h, w)                       
  MSF: 
    ENABLE      : false                                   # multi-scale and flip evaluation  
    FLIP        : true                                    # use flip in evaluation  
    SCALES      : [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]       # scales used in MSF evaluation                

TEST:
  MODEL_PATH    : '/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/DDR-NET_output/ddrnet_23slim_city.pth'    # trained model file path
  FILE          : 'assests/cityscapes'                    # filename or foldername 
  IMAGE_SIZE    : [1024, 1024]                            # inference image size in (h, w)
  OVERLAY       : true                                    # save the overlay result (image_alpha+label_alpha)

but when I start training I get the following error:

Traceback (most recent call last):
  File "/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/DDR-NET/semantic-segmentation/tools/train.py", line 128, in <module>
    main(cfg, gpu, save_dir)
  File "/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/DDR-NET/semantic-segmentation/tools/train.py", line 42, in main
    model.init_pretrained(model_cfg['PRETRAINED'])
  File "/home/banana/Docs/VScode/Python/RSM_projects/Auto_Pilot/Kromka_Semantic_test/DDR-NET/semantic-segmentation/semseg/models/ddrnet.py", line 191, in init_pretrained
    self.load_state_dict(torch.load(pretrained, map_location='cpu')['model'], strict=False)
KeyError: 'model'