Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| # let us run this cell only if CUDA is available | |
| # We will use ``torch.device`` objects to move tensors in and out of GPU | |
| if torch.cuda.is_available(): | |
| x = torch.randn(1) | |
| device = torch.device("cuda") # a CUDA device object | |
| y = torch.ones_like(x, device=device) # directly create a tensor on GPU | |
| x = x.to(device) # or just use strings ``.to("cuda")`` | |
| z = x + y | |
| print(z) | |
| print(z.to("cpu", torch.double)) # ``.to`` can also change dtype together! | |
| class YoloLayer(nn.Module): | |
| def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1): | |
| super(YoloLayer, self).__init__() | |
| self.anchor_mask = anchor_mask | |
| self.num_classes = num_classes | |
| self.anchors = anchors | |
| self.num_anchors = num_anchors | |
| self.anchor_step = len(anchors)/num_anchors | |
| self.coord_scale = 1 | |
| self.noobject_scale = 1 | |
| self.object_scale = 5 | |
| self.class_scale = 1 | |
| self.thresh = 0.6 | |
| self.stride = 32 | |
| self.seen = 0 | |
| def forward(self, output, nms_thresh): | |
| self.thresh = nms_thresh | |
| masked_anchors = [] | |
| for m in self.anchor_mask: | |
| masked_anchors += self.anchors[m*self.anchor_step:(m+1)*self.anchor_step] | |
| masked_anchors = [anchor/self.stride for anchor in masked_anchors] | |
| boxes = get_region_boxes(output.data, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask)) | |
| return boxes | |
| class Upsample(nn.Module): | |
| def __init__(self, stride=2): | |
| super(Upsample, self).__init__() | |
| self.stride = stride | |
| def forward(self, x): | |
| stride = self.stride | |
| assert(x.data.dim() == 4) | |
| B = x.data.size(0) | |
| C = x.data.size(1) | |
| H = x.data.size(2) | |
| W = x.data.size(3) | |
| ws = stride | |
| hs = stride | |
| x = x.view(B, C, H, 1, W, 1).expand(B, C, H, stride, W, stride).contiguous().view(B, C, H*stride, W*stride) | |
| return x | |
| #for route and shortcut | |
| class EmptyModule(nn.Module): | |
| def __init__(self): | |
| super(EmptyModule, self).__init__() | |
| def forward(self, x): | |
| return x | |
| # support route shortcut | |
| class Darknet(nn.Module): | |
| def __init__(self, cfgfile): | |
| super(Darknet, self).__init__() | |
| self.blocks = parse_cfg(cfgfile) | |
| self.models = self.create_network(self.blocks) # merge conv, bn,leaky | |
| self.loss = self.models[len(self.models)-1] | |
| self.width = int(self.blocks[0]['width']) | |
| self.height = int(self.blocks[0]['height']) | |
| self.header = torch.IntTensor([0,0,0,0]) | |
| self.seen = 0 | |
| def forward(self, x, nms_thresh): | |
| ind = -2 | |
| self.loss = None | |
| outputs = dict() | |
| out_boxes = [] | |
| for block in self.blocks: | |
| ind = ind + 1 | |
| if block['type'] == 'net': | |
| continue | |
| elif block['type'] in ['convolutional', 'upsample']: | |
| x = self.models[ind](x) | |
| outputs[ind] = x | |
| elif block['type'] == 'route': | |
| layers = block['layers'].split(',') | |
| layers = [int(i) if int(i) > 0 else int(i)+ind for i in layers] | |
| if len(layers) == 1: | |
| x = outputs[layers[0]] | |
| outputs[ind] = x | |
| elif len(layers) == 2: | |
| x1 = outputs[layers[0]] | |
| x2 = outputs[layers[1]] | |
| x = torch.cat((x1,x2),1) | |
| outputs[ind] = x | |
| elif block['type'] == 'shortcut': | |
| from_layer = int(block['from']) | |
| activation = block['activation'] | |
| from_layer = from_layer if from_layer > 0 else from_layer + ind | |
| x1 = outputs[from_layer] | |
| x2 = outputs[ind-1] | |
| x = x1 + x2 | |
| outputs[ind] = x | |
| elif block['type'] == 'yolo': | |
| boxes = self.models[ind](x, nms_thresh) | |
| out_boxes.append(boxes) | |
| else: | |
| print('unknown type %s' % (block['type'])) | |
| return out_boxes | |
| def print_network(self): | |
| print_cfg(self.blocks) | |
| def create_network(self, blocks): | |
| models = nn.ModuleList() | |
| prev_filters = 3 | |
| out_filters =[] | |
| prev_stride = 1 | |
| out_strides = [] | |
| conv_id = 0 | |
| for block in blocks: | |
| if block['type'] == 'net': | |
| prev_filters = int(block['channels']) | |
| continue | |
| elif block['type'] == 'convolutional': | |
| conv_id = conv_id + 1 | |
| batch_normalize = int(block['batch_normalize']) | |
| filters = int(block['filters']) | |
| kernel_size = int(block['size']) | |
| stride = int(block['stride']) | |
| is_pad = int(block['pad']) | |
| pad = (kernel_size-1)//2 if is_pad else 0 | |
| activation = block['activation'] | |
| model = nn.Sequential() | |
| if batch_normalize: | |
| model.add_module('conv{0}'.format(conv_id), nn.Conv2d(prev_filters, filters, kernel_size, stride, pad, bias=False)) | |
| model.add_module('bn{0}'.format(conv_id), nn.BatchNorm2d(filters)) | |
| else: | |
| model.add_module('conv{0}'.format(conv_id), nn.Conv2d(prev_filters, filters, kernel_size, stride, pad)) | |
| if activation == 'leaky': | |
| model.add_module('leaky{0}'.format(conv_id), nn.LeakyReLU(0.1, inplace=True)) | |
| prev_filters = filters | |
| out_filters.append(prev_filters) | |
| prev_stride = stride * prev_stride | |
| out_strides.append(prev_stride) | |
| models.append(model) | |
| elif block['type'] == 'upsample': | |
| stride = int(block['stride']) | |
| out_filters.append(prev_filters) | |
| prev_stride = prev_stride // stride | |
| out_strides.append(prev_stride) | |
| models.append(Upsample(stride)) | |
| elif block['type'] == 'route': | |
| layers = block['layers'].split(',') | |
| ind = len(models) | |
| layers = [int(i) if int(i) > 0 else int(i)+ind for i in layers] | |
| if len(layers) == 1: | |
| prev_filters = out_filters[layers[0]] | |
| prev_stride = out_strides[layers[0]] | |
| elif len(layers) == 2: | |
| assert(layers[0] == ind - 1) | |
| prev_filters = out_filters[layers[0]] + out_filters[layers[1]] | |
| prev_stride = out_strides[layers[0]] | |
| out_filters.append(prev_filters) | |
| out_strides.append(prev_stride) | |
| models.append(EmptyModule()) | |
| elif block['type'] == 'shortcut': | |
| ind = len(models) | |
| prev_filters = out_filters[ind-1] | |
| out_filters.append(prev_filters) | |
| prev_stride = out_strides[ind-1] | |
| out_strides.append(prev_stride) | |
| models.append(EmptyModule()) | |
| elif block['type'] == 'yolo': | |
| yolo_layer = YoloLayer() | |
| anchors = block['anchors'].split(',') | |
| anchor_mask = block['mask'].split(',') | |
| yolo_layer.anchor_mask = [int(i) for i in anchor_mask] | |
| yolo_layer.anchors = [float(i) for i in anchors] | |
| yolo_layer.num_classes = int(block['classes']) | |
| yolo_layer.num_anchors = int(block['num']) | |
| yolo_layer.anchor_step = len(yolo_layer.anchors)//yolo_layer.num_anchors | |
| yolo_layer.stride = prev_stride | |
| out_filters.append(prev_filters) | |
| out_strides.append(prev_stride) | |
| models.append(yolo_layer) | |
| else: | |
| print('unknown type %s' % (block['type'])) | |
| return models | |
| def load_weights(self, weightfile): | |
| print() | |
| fp = open(weightfile, 'rb') | |
| header = np.fromfile(fp, count=5, dtype=np.int32) | |
| self.header = torch.from_numpy(header) | |
| self.seen = self.header[3] | |
| buf = np.fromfile(fp, dtype = np.float32) | |
| fp.close() | |
| start = 0 | |
| ind = -2 | |
| counter = 3 | |
| for block in self.blocks: | |
| if start >= buf.size: | |
| break | |
| ind = ind + 1 | |
| if block['type'] == 'net': | |
| continue | |
| elif block['type'] == 'convolutional': | |
| model = self.models[ind] | |
| batch_normalize = int(block['batch_normalize']) | |
| if batch_normalize: | |
| start = load_conv_bn(buf, start, model[0], model[1]) | |
| else: | |
| start = load_conv(buf, start, model[0]) | |
| elif block['type'] == 'upsample': | |
| pass | |
| elif block['type'] == 'route': | |
| pass | |
| elif block['type'] == 'shortcut': | |
| pass | |
| elif block['type'] == 'yolo': | |
| pass | |
| else: | |
| print('unknown type %s' % (block['type'])) | |
| percent_comp = (counter / len(self.blocks)) * 100 | |
| print('Loading weights. Please Wait...{:.2f}% Complete'.format(percent_comp), end = '\r', flush = True) | |
| counter += 1 | |
| def convert2cpu(gpu_matrix): | |
| return torch.FloatTensor(gpu_matrix.size()).copy_(gpu_matrix) | |
| def convert2cpu_long(gpu_matrix): | |
| return torch.LongTensor(gpu_matrix.size()).copy_(gpu_matrix) | |
| def get_region_boxes(output, conf_thresh, num_classes, anchors, num_anchors, only_objectness = 1, validation = False): | |
| anchor_step = len(anchors)//num_anchors | |
| if output.dim() == 3: | |
| output = output.unsqueeze(0) | |
| batch = output.size(0) | |
| assert(output.size(1) == (5+num_classes)*num_anchors) | |
| h = output.size(2) | |
| w = output.size(3) | |
| all_boxes = [] | |
| output = output.view(batch*num_anchors, 5+num_classes, h*w).transpose(0,1).contiguous().view(5+num_classes, batch*num_anchors*h*w) | |
| grid_x = torch.linspace(0, w-1, w).repeat(h,1).repeat(batch*num_anchors, 1, 1).view(batch*num_anchors*h*w).type_as(output) #cuda() | |
| grid_y = torch.linspace(0, h-1, h).repeat(w,1).t().repeat(batch*num_anchors, 1, 1).view(batch*num_anchors*h*w).type_as(output) #cuda() | |
| xs = torch.sigmoid(output[0]) + grid_x | |
| ys = torch.sigmoid(output[1]) + grid_y | |
| anchor_w = torch.Tensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([0])) | |
| anchor_h = torch.Tensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([1])) | |
| anchor_w = anchor_w.repeat(batch, 1).repeat(1, 1, h*w).view(batch*num_anchors*h*w).type_as(output) #cuda() | |
| anchor_h = anchor_h.repeat(batch, 1).repeat(1, 1, h*w).view(batch*num_anchors*h*w).type_as(output) #cuda() | |
| ws = torch.exp(output[2]) * anchor_w | |
| hs = torch.exp(output[3]) * anchor_h | |
| det_confs = torch.sigmoid(output[4]) | |
| cls_confs = torch.nn.Softmax(dim=1)(output[5:5+num_classes].transpose(0,1)).detach() | |
| cls_max_confs, cls_max_ids = torch.max(cls_confs, 1) | |
| cls_max_confs = cls_max_confs.view(-1) | |
| cls_max_ids = cls_max_ids.view(-1) | |
| sz_hw = h*w | |
| sz_hwa = sz_hw*num_anchors | |
| det_confs = convert2cpu(det_confs) | |
| cls_max_confs = convert2cpu(cls_max_confs) | |
| cls_max_ids = convert2cpu_long(cls_max_ids) | |
| xs = convert2cpu(xs) | |
| ys = convert2cpu(ys) | |
| ws = convert2cpu(ws) | |
| hs = convert2cpu(hs) | |
| if validation: | |
| cls_confs = convert2cpu(cls_confs.view(-1, num_classes)) | |
| for b in range(batch): | |
| boxes = [] | |
| for cy in range(h): | |
| for cx in range(w): | |
| for i in range(num_anchors): | |
| ind = b*sz_hwa + i*sz_hw + cy*w + cx | |
| det_conf = det_confs[ind] | |
| if only_objectness: | |
| conf = det_confs[ind] | |
| else: | |
| conf = det_confs[ind] * cls_max_confs[ind] | |
| if conf > conf_thresh: | |
| bcx = xs[ind] | |
| bcy = ys[ind] | |
| bw = ws[ind] | |
| bh = hs[ind] | |
| cls_max_conf = cls_max_confs[ind] | |
| cls_max_id = cls_max_ids[ind] | |
| box = [bcx/w, bcy/h, bw/w, bh/h, det_conf, cls_max_conf, cls_max_id] | |
| if (not only_objectness) and validation: | |
| for c in range(num_classes): | |
| tmp_conf = cls_confs[ind][c] | |
| if c != cls_max_id and det_confs[ind]*tmp_conf > conf_thresh: | |
| box.append(tmp_conf) | |
| box.append(c) | |
| boxes.append(box) | |
| all_boxes.append(boxes) | |
| return all_boxes | |
| def parse_cfg(cfgfile): | |
| blocks = [] | |
| fp = open(cfgfile, 'r') | |
| block = None | |
| line = fp.readline() | |
| while line != '': | |
| line = line.rstrip() | |
| if line == '' or line[0] == '#': | |
| line = fp.readline() | |
| continue | |
| elif line[0] == '[': | |
| if block: | |
| blocks.append(block) | |
| block = dict() | |
| block['type'] = line.lstrip('[').rstrip(']') | |
| # set default value | |
| if block['type'] == 'convolutional': | |
| block['batch_normalize'] = 0 | |
| else: | |
| key,value = line.split('=') | |
| key = key.strip() | |
| if key == 'type': | |
| key = '_type' | |
| value = value.strip() | |
| block[key] = value | |
| line = fp.readline() | |
| if block: | |
| blocks.append(block) | |
| fp.close() | |
| return blocks | |
| def print_cfg(blocks): | |
| print('layer filters size input output') | |
| prev_width = 416 | |
| prev_height = 416 | |
| prev_filters = 3 | |
| out_filters =[] | |
| out_widths =[] | |
| out_heights =[] | |
| ind = -2 | |
| for block in blocks: | |
| ind = ind + 1 | |
| if block['type'] == 'net': | |
| prev_width = int(block['width']) | |
| prev_height = int(block['height']) | |
| continue | |
| elif block['type'] == 'convolutional': | |
| filters = int(block['filters']) | |
| kernel_size = int(block['size']) | |
| stride = int(block['stride']) | |
| is_pad = int(block['pad']) | |
| pad = (kernel_size-1)//2 if is_pad else 0 | |
| width = (prev_width + 2*pad - kernel_size)//stride + 1 | |
| height = (prev_height + 2*pad - kernel_size)//stride + 1 | |
| print('%5d %-6s %4d %d x %d / %d %3d x %3d x%4d -> %3d x %3d x%4d' % (ind, 'conv', filters, kernel_size, kernel_size, stride, prev_width, prev_height, prev_filters, width, height, filters)) | |
| prev_width = width | |
| prev_height = height | |
| prev_filters = filters | |
| out_widths.append(prev_width) | |
| out_heights.append(prev_height) | |
| out_filters.append(prev_filters) | |
| elif block['type'] == 'upsample': | |
| stride = int(block['stride']) | |
| filters = prev_filters | |
| width = prev_width*stride | |
| height = prev_height*stride | |
| print('%5d %-6s * %d %3d x %3d x%4d -> %3d x %3d x%4d' % (ind, 'upsample', stride, prev_width, prev_height, prev_filters, width, height, filters)) | |
| prev_width = width | |
| prev_height = height | |
| prev_filters = filters | |
| out_widths.append(prev_width) | |
| out_heights.append(prev_height) | |
| out_filters.append(prev_filters) | |
| elif block['type'] == 'route': | |
| layers = block['layers'].split(',') | |
| layers = [int(i) if int(i) > 0 else int(i)+ind for i in layers] | |
| if len(layers) == 1: | |
| print('%5d %-6s %d' % (ind, 'route', layers[0])) | |
| prev_width = out_widths[layers[0]] | |
| prev_height = out_heights[layers[0]] | |
| prev_filters = out_filters[layers[0]] | |
| elif len(layers) == 2: | |
| print('%5d %-6s %d %d' % (ind, 'route', layers[0], layers[1])) | |
| prev_width = out_widths[layers[0]] | |
| prev_height = out_heights[layers[0]] | |
| assert(prev_width == out_widths[layers[1]]) | |
| assert(prev_height == out_heights[layers[1]]) | |
| prev_filters = out_filters[layers[0]] + out_filters[layers[1]] | |
| out_widths.append(prev_width) | |
| out_heights.append(prev_height) | |
| out_filters.append(prev_filters) | |
| elif block['type'] in ['region', 'yolo']: | |
| print('%5d %-6s' % (ind, 'detection')) | |
| out_widths.append(prev_width) | |
| out_heights.append(prev_height) | |
| out_filters.append(prev_filters) | |
| elif block['type'] == 'shortcut': | |
| from_id = int(block['from']) | |
| from_id = from_id if from_id > 0 else from_id+ind | |
| print('%5d %-6s %d' % (ind, 'shortcut', from_id)) | |
| prev_width = out_widths[from_id] | |
| prev_height = out_heights[from_id] | |
| prev_filters = out_filters[from_id] | |
| out_widths.append(prev_width) | |
| out_heights.append(prev_height) | |
| out_filters.append(prev_filters) | |
| else: | |
| print('unknown type %s' % (block['type'])) | |
| def load_conv(buf, start, conv_model): | |
| num_w = conv_model.weight.numel() | |
| num_b = conv_model.bias.numel() | |
| conv_model.bias.data.copy_(torch.from_numpy(buf[start:start+num_b])); start = start + num_b | |
| conv_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_w]).view_as(conv_model.weight.data)); start = start + num_w | |
| return start | |
| def load_conv_bn(buf, start, conv_model, bn_model): | |
| num_w = conv_model.weight.numel() | |
| num_b = bn_model.bias.numel() | |
| bn_model.bias.data.copy_(torch.from_numpy(buf[start:start+num_b])); start = start + num_b | |
| bn_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_b])); start = start + num_b | |
| bn_model.running_mean.copy_(torch.from_numpy(buf[start:start+num_b])); start = start + num_b | |
| bn_model.running_var.copy_(torch.from_numpy(buf[start:start+num_b])); start = start + num_b | |
| conv_model.weight.data.copy_(torch.from_numpy(buf[start:start+num_w]).view_as(conv_model.weight.data)); start = start + num_w | |
| return start | |