def get_transform(train, advanced=False):
transforms = []
if train:
# during training, randomly flip the training images
# and ground-truth for data augmentation
transforms_list = [
# bboxAug.RandomChoice(),
# ---------两者取其一--------------------
bboxAug.RandomHorizontalFlip(),
# bboxAug.RandomTranslate(), # 如果有mask也需相应修改
# bboxAug.RandomRotate(3),
bboxAug.RandomBrightness(),
bboxAug.RandomSaturation(),
bboxAug.RandomHue(),
bboxAug.RandomBlur(),
# ---------两者取其一--------------------
# bboxAug.Augment(advanced),
# ---------两者取其一--------------------
]
transforms.extend(transforms_list)
transforms.append(bboxAug.ToTensor())
return T.Compose(transforms)
def __init__(self,trainDP=None,testDP=None,model_name="resnet18",num_features=None,
pretrained=False,dropRate=0.0, usize=256,isTrain=False,useRCNN=False,
basePath="./",save_model = "model.pt",summaryPath="yolov1_resnet50_416",
epochs = 100,print_freq=1,resize:tuple = (224,224),
mulScale=False,advanced=False,batch_size=2,num_anchors=6,lr=2e-3,
typeOfData="PennFudanDataset",
threshold_conf=0.5,threshold_cls=0.5, # # 0.05,0.5
conf_thres=0.7,nms_thres=0.4, # 0.8,0.4
filter_labels = [],classes=[]):
super(CascadeSSD,self).__init__()
self.batch_size = batch_size
self.epochs = epochs
self.print_freq = print_freq
self.isTrain = isTrain
self.mulScale = mulScale
self.classes = classes
num_classes = len(self.classes)
num_anchors = 6
self.useRCNN = useRCNN
# seed = 100
seed = int(time.time() * 1000)
self.use_cuda = torch.cuda.is_available()
self.device = torch.device("cuda" if self.use_cuda else "cpu")
torch.manual_seed(seed)
kwargs = {'num_workers': 5, 'pin_memory': True} if self.use_cuda else {}
if self.isTrain:
if typeOfData=="PennFudanDataset":
Data = datasets.PennFudanDataset
elif typeOfData=="PascalVOCDataset":
Data = datasets.PascalVOCDataset
else:
pass
train_dataset = Data(trainDP,
transforms=bboxAug.Compose([
# bboxAug.RandomChoice(),
bboxAug.Pad(), bboxAug.Resize(resize, mulScale),
# *random.choice([
# [bboxAug.Pad(), bboxAug.Resize(resize, mulScale)],
# [bboxAug.Resize2(resize, mulScale)]
# ]),
# ---------两者取其一--------------------
# bboxAug.RandomHorizontalFlip(),
# bboxAug.RandomTranslate(),
# # bboxAug.RandomRotate(3),
# bboxAug.RandomBrightness(),
# bboxAug.RandomSaturation(),
# bboxAug.RandomHue(),
# bboxAug.RandomBlur(),
# bboxAug.Augment(advanced),
# -------------------------------
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]),classes=classes)
test_dataset = datasets.ValidDataset(testDP,
transforms=bboxAug.Compose([
bboxAug.Pad(), bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]))
self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
collate_fn=collate_fn, **kwargs)
self.test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False,
collate_fn=collate_fn, **kwargs)
else:
test_dataset = datasets.ValidDataset(testDP,
transforms=bboxAug.Compose([
bboxAug.Pad(), bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]))
self.test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False,
collate_fn=collate_fn, **kwargs)
self.loss_func = cascadessdLoss.CascadeSSDLoss(self.device,num_anchors,num_classes,threshold_conf,
threshold_cls,conf_thres,nms_thres,filter_labels,self.mulScale,self.useRCNN)
self.network = cascadeSSDNet.CascadeSSDNet(num_classes,num_anchors,model_name,num_features,pretrained,dropRate,usize)
# self.network.apply(cascadeSSDNet.weights_init)
self.network.fpn.apply(cascadeSSDNet.weights_init) # backbone 不使用
self.network.rpnNet.apply(cascadeSSDNet.weights_init2)
self.network.rcnnNet.apply(cascadeSSDNet.weights_init2)
if self.use_cuda:
self.network.to(self.device)
if not os.path.exists(basePath):
os.makedirs(basePath)
self.save_model = os.path.join(basePath,save_model)
if os.path.exists(self.save_model):
self.network.load_state_dict(torch.load(self.save_model))
# self.network.load_state_dict(torch.load(self.save_model,map_location=torch.device('cpu')))
# optimizer
base_params = list(
map(id, self.network.backbone.parameters())
)
# logits_params = filter(lambda p: id(p) not in base_params, self.network.parameters())
params = [
{"params": self.network.fpn.parameters(), "lr": lr}, # 1e-3
{"params": self.network.rpnNet.parameters(), "lr": lr}, # 1e-3
{"params": self.network.backbone.parameters(), "lr": lr / 10}, # 1e-4
]
# self.optimizer = torch.optim.Adam(self.network.parameters(), lr=lr, weight_decay=4e-05)
# self.optimizer = torch.optim.Adam(params, weight_decay=4e-05)
self.optimizer = optimizer.RAdam(params, weight_decay=4e-05)
if self.useRCNN:
self.optimizer_rcnn = torch.optim.Adam(self.network.rcnnNet.parameters(),lr=lr/10, weight_decay=4e-05)
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=20, gamma=0.1)
if self.useRCNN:
self.lr_scheduler_rcnn = torch.optim.lr_scheduler.StepLR(self.optimizer_rcnn, step_size=20, gamma=0.1)
self.writer = SummaryWriter(os.path.join(basePath,summaryPath))
self.history = History()
def __init__(self,cfg):
super(FasterRCNN,self).__init__()
self.batch_size = cfg["work"]["train"]["batch_size"]
self.epochs = cfg["work"]["train"]["epochs"]
self.print_freq = cfg["work"]["train"]["print_freq"]
self.isTrain = cfg["work"]["train"]["isTrain"]
self.classes = cfg["work"]["train"]["classes"]
self.resize = cfg["work"]["train"]["resize"]
num_classes = len(self.classes)
self.train_method = cfg["work"]["train"]["train_method"]
typeOfData = cfg["work"]["dataset"]["typeOfData"]
trainDP = cfg["work"]["dataset"]["trainDataPath"]
testDP = cfg["work"]["dataset"]["testDataPath"]
predDP = cfg["work"]["dataset"]["predDataPath"]
advanced = cfg["work"]["train"]["advanced"]
useImgaug = cfg["work"]["train"]["useImgaug"]
resize = cfg["work"]["train"]["resize"]
use_FPN = cfg["network"]["FPN"]["use_FPN"]
basePath = cfg["work"]["save"]["basePath"]
save_model = cfg["work"]["save"]["save_model"]
summaryPath = cfg["work"]["save"]["summaryPath"]
# lr = cfg["work"]["train"]["lr"]
seed = 100
# seed = int(time.time() * 1000)
self.use_cuda = torch.cuda.is_available()
self.device = torch.device("cuda" if self.use_cuda else "cpu")
torch.manual_seed(seed)
kwargs = {'num_workers': 5, 'pin_memory': True} if self.use_cuda else {}
if useImgaug:
train_transforms = bboxAug.Compose([
bboxAug.Augment(advanced=advanced),
bboxAug.Pad(), bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
])
else:
train_transforms = bboxAug.Compose([
bboxAug.RandomHorizontalFlip(),
bboxAug.RandomBrightness(),
bboxAug.RandomBlur(),
bboxAug.RandomSaturation(),
bboxAug.RandomHue(),
bboxAug.RandomRotate(3),
bboxAug.RandomTranslate(),
bboxAug.Pad(), bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
])
test_transforms = bboxAug.Compose([
bboxAug.Pad(), bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
])
if self.isTrain:
if typeOfData=="PennFudanDataset":
Data = datasets.PennFudanDataset
elif typeOfData=="PascalVOCDataset":
Data = datasets.PascalVOCDataset
else:
pass
train_dataset = Data(trainDP,transforms=train_transforms,classes=self.classes)
if testDP is not None:
test_dataset = Data(testDP,transforms=test_transforms,classes=self.classes)
else:
test_dataset = Data(trainDP,test_transforms, classes=self.classes)
num_datas = len(train_dataset)
num_train = int(0.8*num_datas)
indices = torch.randperm(num_datas).tolist()
train_dataset = torch.utils.data.Subset(train_dataset, indices[:num_train])
test_dataset = torch.utils.data.Subset(test_dataset, indices[num_train:])
self.train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True,
collate_fn=collate_fn, **kwargs)
self.test_loader = DataLoader(test_dataset, batch_size=self.batch_size, shuffle=False,
collate_fn=collate_fn, **kwargs)
if predDP is not None:
pred_dataset = datasets.ValidDataset(predDP,transforms=test_transforms)
self.pred_loader = DataLoader(pred_dataset, batch_size=self.batch_size, shuffle=False,
collate_fn=collate_fn, **kwargs)
else:
pred_dataset = datasets.ValidDataset(predDP,transforms=test_transforms)
self.pred_loader = DataLoader(pred_dataset, batch_size=self.batch_size, shuffle=False,
collate_fn=collate_fn, **kwargs)
self.loss_func = fastrcnnLoss.RPNLoss(cfg, self.device)
self.loss_func_rcnn = fastrcnnLoss.RCNNLoss(cfg, self.device)
self.network = net.FasterRCNN(cfg)
if use_FPN:self.network.fpn.apply(net.weights_init_fpn) # backbone 不使用
self.network.rpn.apply(net.weights_init_rpn)
self.network.rcnn.apply(net.weights_init_rpn)
if self.use_cuda:
self.network.to(self.device)
if not os.path.exists(basePath):
os.makedirs(basePath)
self.save_model = os.path.join(basePath,save_model)
if os.path.exists(self.save_model):
self.network.load_state_dict(torch.load(self.save_model))
# self.network.load_state_dict(torch.load(self.save_model,map_location=torch.device('cpu')))
self.optimizer = optimizer.build_optimizer(self.network,cfg)
self.lr_scheduler = optimizer.build_lr_scheduler(self.optimizer,cfg)
self.writer = SummaryWriter(os.path.join(basePath,summaryPath))
def __init__(
self,
trainDP=None,
testDP=None,
predDP=None,
model_name="resnet18",
num_features=None,
pretrained=False,
dropRate=0.0,
usize=256,
isTrain=False,
freeze_at=1,
basePath="./",
save_model="model.pt",
summaryPath="yolov1_resnet50_416",
epochs=100,
print_freq=1,
resize: tuple = (224, 224),
useFocal=True,
method=1,
mulScale=False,
advanced=False,
batch_size=2,
num_anchors=2,
lr=5e-4,
train_method=1,
# num_classes=20,
typeOfData="PennFudanDataset",
threshold_conf=0.5,
threshold_cls=0.5, # # 0.05,0.5
conf_thres=0.5,
nms_thres=0.4, # 0.8,0.4
filter_labels=[],
classes=[],
version="v1"): # "v1" v2 v3
super(YOLO, self).__init__()
self.batch_size = batch_size
self.epochs = epochs
self.print_freq = print_freq
self.isTrain = isTrain
self.mulScale = mulScale
self.classes = classes
num_classes = len(self.classes)
self.train_method = train_method
# seed = 100
seed = int(time.time() * 1000)
self.use_cuda = torch.cuda.is_available()
self.device = torch.device("cuda" if self.use_cuda else "cpu")
torch.manual_seed(seed)
kwargs = {
'num_workers': 5,
'pin_memory': True
} if self.use_cuda else {}
if self.isTrain:
if typeOfData == "PennFudanDataset":
Data = datasets.PennFudanDataset
elif typeOfData == "PascalVOCDataset":
Data = datasets.PascalVOCDataset
else:
pass
train_dataset = Data(
trainDP,
transforms=bboxAug.Compose([
# bboxAug.RandomChoice(),
bboxAug.Pad(),
bboxAug.Resize(resize, mulScale),
# *random.choice([
# [bboxAug.Pad(), bboxAug.Resize(resize, mulScale)],
# [bboxAug.Resize2(resize, mulScale)]
# ]),
# ---------两者取其一--------------------
# bboxAug.RandomHorizontalFlip(),
# bboxAug.RandomTranslate(),
# # bboxAug.RandomRotate(3),
# bboxAug.RandomBrightness(),
# bboxAug.RandomSaturation(),
# bboxAug.RandomHue(),
# bboxAug.RandomBlur(),
# bboxAug.Augment(advanced),
# -------------------------------
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]),
classes=classes)
if testDP is not None:
test_dataset = Data(
testDP,
transforms=bboxAug.Compose([
bboxAug.Pad(),
bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]),
classes=classes)
else:
test_dataset = Data(
trainDP,
transforms=bboxAug.Compose([
bboxAug.Pad(),
bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]),
classes=classes)
num_datas = len(train_dataset)
num_train = int(0.8 * num_datas)
indices = torch.randperm(num_datas).tolist()
train_dataset = torch.utils.data.Subset(
train_dataset, indices[:num_train])
test_dataset = torch.utils.data.Subset(test_dataset,
indices[num_train:])
self.train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=collate_fn,
**kwargs)
self.test_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
collate_fn=collate_fn,
**kwargs)
if predDP is not None:
pred_dataset = datasets.ValidDataset(
predDP,
transforms=bboxAug.Compose([
bboxAug.Pad(),
bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]))
self.pred_loader = DataLoader(pred_dataset,
batch_size=batch_size,
shuffle=False,
collate_fn=collate_fn,
**kwargs)
else:
pred_dataset = datasets.ValidDataset(
predDP,
transforms=bboxAug.Compose([
bboxAug.Pad(),
bboxAug.Resize(resize, False),
bboxAug.ToTensor(), # PIL --> tensor
bboxAug.Normalize() # tensor --> tensor
]))
self.pred_loader = DataLoader(pred_dataset,
batch_size=batch_size,
shuffle=False,
collate_fn=collate_fn,
**kwargs)
if version == "v1":
self.loss_func = yoloLoss.YOLOv1Loss(self.device, num_anchors,
num_classes, threshold_conf,
threshold_cls, conf_thres,
nms_thres, filter_labels,
self.mulScale, useFocal)
else:
self.loss_func = yoloLoss.YOLOv2Loss(self.device, num_anchors,
num_classes, threshold_conf,
threshold_cls, conf_thres,
nms_thres, filter_labels,
self.mulScale, useFocal,
method)
self.network = yoloNet.YOLOV1Net(num_classes,
self.loss_func.num_anchors,
model_name, num_features, pretrained,
dropRate, usize, freeze_at)
# self.network.apply(yoloNet.weights_init)
self.network.fpn.apply(yoloNet.weights_init) # backbone 不使用
self.network.net.apply(yoloNet.weights_init2)
if self.use_cuda:
self.network.to(self.device)
if not os.path.exists(basePath):
os.makedirs(basePath)
self.save_model = os.path.join(basePath, save_model)
if os.path.exists(self.save_model):
self.network.load_state_dict(torch.load(self.save_model))
# self.network.load_state_dict(torch.load(self.save_model,map_location=torch.device('cpu')))
"""
# optimizer
base_params = list(
map(id, self.network.backbone.parameters())
)
logits_params = filter(lambda p: id(p) not in base_params, self.network.parameters())
params = [
{"params": logits_params, "lr": lr}, # 1e-3
{"params": self.network.backbone.parameters(), "lr": lr / 10}, # 1e-4
]
# self.optimizer = torch.optim.Adam(self.network.parameters(), lr=lr, weight_decay=4e-05)
# self.optimizer = torch.optim.Adam(params, weight_decay=4e-05)
self.optimizer = optimizer.RAdam(params, weight_decay=4e-05)
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=20, gamma=0.1)
"""
self.optimizer = optimizer.build_optimizer(self.network,
lr,
clip_gradients=True)
self.lr_scheduler = optimizer.build_lr_scheduler(self.optimizer)
self.writer = SummaryWriter(os.path.join(basePath, summaryPath))
self.history = History()