MCPcopy Create free account
hub / github.com/Modulus-Labs/RockyBot / main

Function main

pytorch-model/classification_eval.py:105–161  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

103
104
105def main():
106 # --- Args ---
107 args = opts.get_eval_args()
108 print("\n" + "-" * 30 + " Args " + "-" * 30)
109 for k, v in vars(args).items():
110 print(f"{k}: {v}")
111 print()
112
113 # --- Model and viz save dir ---
114 model_save_dir = constants.get_model_dir(args.dataset, args.model_type, args.model_name)
115 viz_save_dir = constants.get_viz_dir(args.dataset, args.model_type, args.model_name)
116 if not os.path.isdir(model_save_dir):
117 raise RuntimeError(f"Error: {model_save_dir} does not exist! Exiting...\n")
118 if not os.path.isdir(viz_save_dir):
119 raise RuntimeError(f"Error: {viz_save_dir} does not exist! Exiting...\n")
120
121 # --- Get model weight path ---
122 model_filenames = sorted(list(os.path.basename(x) for x in glob.glob(os.path.join(model_save_dir, "*.pth"))))
123 for model_file in model_filenames:
124 print(model_file)
125 user_choice = input("Please select which model .pth file to load -> ")
126 while user_choice not in model_filenames:
127 user_choice = input (f"Error. Failed to find specified model. Please try again -> ")
128 model_weights_path = os.path.join(model_save_dir, user_choice)
129 print(f"--> Selected {model_weights_path} as the model weights file.\n")
130
131 # --- Setup dataset ---
132 print("--> Setting up dataset...")
133 val_dataset = datasets.DATASETS[args.dataset](mode="val")
134 print("Done!\n")
135
136 # --- Dataloaders ---
137 print("--> Setting up dataloaders...")
138 val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=1)
139 print("Done!\n")
140
141 # --- Setup model ---
142 # TODO(ryancao): Actually pull the ResNet model! ---
143 print("--> Setting up model...")
144 model = models.MODEL_TYPES[args.model_type](val_dataset)
145 model.load_state_dict(torch.load(model_weights_path, map_location=torch.device("cpu")))
146 # torch.cuda.set_device(constants.GPU)
147 # model = model.cuda(constants.GPU)
148 print("Done!\n")
149
150 # --- Loss fn ---
151 print("--> Setting up loss fn...")
152 criterion = nn.CrossEntropyLoss()#.cuda(constants.GPU)
153 print("Done!\n")
154
155 # --- Run eval ---
156 val_avg_loss, val_avg_acc, total_examples = eval_model(model,
157 val_dataloader,
158 criterion,
159 args)
160 print(f"Avg loss: {val_avg_loss} | Avg acc: {val_avg_acc} | "\
161 f"Total number of val examples: {total_examples}")
162

Callers 1

Calls 1

eval_modelFunction · 0.70

Tested by

no test coverage detected