def __init__(self,
device,
input_size=None,
num_classes=20,
trainable=False,
conf_thresh=0.01,
nms_thresh=0.5,
hr=False):
super(myYOLO, self).__init__()
self.device = device
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = 32
self.grid_cell = self.create_grid(input_size)
self.input_size = input_size
self.scale = np.array(
[[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
self.scale_torch = torch.tensor(self.scale.copy(),
device=device).float()
# we use resnet18 as backbone
self.backbone = resnet18(pretrained=True)
# neck
self.SPP = nn.Sequential(
Conv(512, 256, k=1), SPP(),
BottleneckCSP(256 * 4, 512, n=1, shortcut=False))
self.SAM = SAM(512)
self.conv_set = BottleneckCSP(512, 512, n=3, shortcut=False)
self.pred = nn.Conv2d(512, 1 + self.num_classes + 4, 1)
def __init__(self, device, input_size=None, num_classes=20, trainable=False, conf_thresh=0.01, nms_thresh=0.5, hr=False):
super().__init__()
self.device = device
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
# 下采样倍数
self.stride =32
# input_size [h,w] 训练时候resize的宽和高 比如416 x 416
self.grid_cell = self.create_grid(input_size)
self.input_size = input_size
# 相对大小
self.scale = np.array([[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
self.scale_torch = torch.tensor(self.scale.copy(), device=device).float()
# 网络结构
self.backbone = resnet18(pretrained=True)
self.pred = nn.Conv2d(512, 1 + self.num_classes + 4, 1)
def __init__(self, input_size=None, num_class=20, trainable=False):
super(YoloV1, self).__init__()
# 参数
self.num_class = num_class
# 是否是训练阶段
self.trainable = trainable
self.stride = 32
self.input_size = input_size
self.grid_cell = self.create_grid()
self.scale = torch.tensor(
[[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
# YoloV1 network
self.backbone = resnet18(pretrained=True)
self.spp = SPP(512, 512)
self.sam = SAM(512)
self.conv_set = nn.Sequential(
Conv2d(512, 256, kernel_size=1, leaky_relu=True),
Conv2d(256, 512, kernel_size=3, padding=1, leaky_relu=True),
Conv2d(512, 256, kernel_size=1, leaky_relu=True),
Conv2d(256, 512, kernel_size=3, padding=1, leaky_relu=True))
self.pred = Conv2d(512, 1 + self.num_class + 4, kernel_size=1)
n_gpu = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
args.distributed = n_gpu > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='gloo',
init_method='env://')
synchronize()
device = 'cuda'
valid_set = COCODataset("/data/COCO_17/", 'val',
preset_transform(args, train=False))
# backbone = vovnet39(pretrained=False)
backbone = resnet18(pretrained=False)
model = FCOS(args, backbone)
# load weight
model_file = "./training_dir/epoch-1.pt"
chkpt = torch.load(model_file, map_location='cpu') # load checkpoint
model.load_state_dict(chkpt['model'])
print('load weights from ' + model_file)
model = model.to(device)
if args.distributed:
model = nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,