def __init__(self, model_file='put_your_model_file(or files)_name_here'):
# You should
# 1. create the model object
# 2. load your state_dict
# 3. call cuda()
# self.model = ...
#
self.batch_size = 1
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.model1 = get_seg_model(get_config()).to(self.device)
self.model1.load_state_dict(torch.load('../PT_FILES/HRNET_RM_pretrain03.pt', map_location=self.device))
self.model2 = BoundingBox().double().to(self.device)
self.model2.load_state_dict(torch.load('../PT_FILES/bbox_no_pretrain_50.pt', map_location=self.device))
pass
def __init__(self, model_file='put_your_model_file_name_here'):
# You should
# 1. create the model object
# 2. load your state_dict
# 3. call cuda()
# self.model = ...
#
self.batch_size = 1
# set device
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
#### model 1: predict Binary RoadMap ####
self.model1 = get_seg_model(get_config()).to(self.device)
self.model1.load_state_dict(
torch.load('HRNET_RM_labeled_data01.pt', map_location=self.device))
# TODO: self.model1.load_state_dict(torch.load('classification.pth', map_location=self.device))
#### model 2: predict Bounding Boxes ####
self.model2 = BoundingBox().to(self.device)
# TODO: self.model2.load_state_dict(torch.load('classification.pth', map_location=self.device))
pass
import hrnet
import torch
import config.hrnet_config as cfg
if __name__ == "__main__":
net = hrnet.get_seg_model(cfg.hr_config)
img = torch.ones((1, 3, 256, 256))
out = net(img)
print(out.shape)
transform=transform,
extra_info=False
)
labeled_valset = LabeledDataset(
image_folder=image_folder,
annotation_file=annotation_csv,
scene_index=val_index,
transform=transform,
extra_info=False
)
trainloader = torch.utils.data.DataLoader(labeled_trainset, batch_size=2, shuffle=True, num_workers=2, collate_fn=collate_fn)
valloader = torch.utils.data.DataLoader(labeled_valset, batch_size=2, shuffle=True, num_workers=2, collate_fn=collate_fn)
model = get_seg_model(get_config()).to(device)
# for param in model.parameters():
# param.requires_grad = True
criterion = torch.nn.BCELoss()
#param_list = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
[{'params': filter(lambda p: p.requires_grad, model.parameters()),
'lr': 0.0001}],
lr=0.0001,
momentum=0.9,
weight_decay=0.0001,
nesterov=False,
)
best_val_loss = 100