def test_lstm():
class Net(nn.Module):
def __init__(self,
vocab_size=20,
embed_dim=300,
hidden_dim=512,
num_layers=2):
super().__init__()
self.hidden_dim = hidden_dim
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.encoder = nn.LSTM(embed_dim,
hidden_dim,
num_layers=num_layers)
self.decoder = nn.Linear(hidden_dim, vocab_size)
def forward(self, x):
embed = self.embedding(x)
out, hidden = self.encoder(embed)
out = self.decoder(out)
out = out.view(-1, out.size(2))
return out, hidden
inp = torch.zeros((1, 100)).long() # [length, batch_size]
df, df_total = summary(Net(), inp)
def test_multi_inputs():
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1_x = nn.Conv2d(1, 10, kernel_size=5)
self.conv2_x = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop_x = nn.Dropout2d()
self.conv1_y = nn.Conv2d(1, 10, kernel_size=5)
self.conv2_y = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop_y = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x, y):
y = F.relu(F.max_pool2d(self.conv1_y(y), 2))
y = F.relu(F.max_pool2d(self.conv2_drop_y(self.conv2_y(y)), 2))
x = F.relu(F.max_pool2d(self.conv1_x(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop_x(self.conv2_x(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
inp1 = torch.zeros((1, 1, 28, 28))
inp2 = torch.zeros((1, 1, 64, 64))
summary_value, _ = summary(Net(), inp1, inp2)
def __init__(self, cfg):
super(VAE, self).__init__(cfg)
# Build Encoder
self._encoder_construct()
# Build Decoder
self._decoder_construct()
# Model Summary
model_summary, _ = summary(self, torch.zeros((2, self.cfg.input_dim)))
logging.info(f"====Model Summary====\n{model_summary}\n")
def __init__(self, cfg):
super(VAE, self).__init__(cfg)
self.latent1_dim = 20
self.in_dim = 32
self.layer1_dim = 64
self.layer2_dim = 128
self.latent2_dim = 20
# Build Encoder
self._encoder_construct()
# Build Decoder
self._decoder_construct()
# Model Summary
model_summary, _ = summary(self, torch.zeros((2, self.cfg.input_dim)))
logging.info(f"====Model Summary====\n{model_summary}\n")
def test_conv():
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
summary_value, _ = summary(Net(), torch.zeros((1, 1, 28, 28)))
def test_resnet_50():
model = torchvision.models.resnet50()
summary_value, _ = summary(model, torch.zeros((4, 3, 224, 224)))
x = self.conv5(x)
x = self.last(x)
return x
def resnet18(in_c, n_cls):
return ResNet(in_c, n_cls, [2, 2, 2, 2], BasicBlock)
def resnet34(in_c, n_cls):
return ResNet(in_c, n_cls, [3, 4, 6, 3], BasicBlock)
def resnet50(in_c, n_cls):
return ResNet(in_c, n_cls, [3, 4, 6, 3])
def resnet101(in_c, n_cls):
return ResNet(in_c, n_cls, [3, 4, 23, 3])
def resnet152(in_c, n_cls):
return ResNet(in_c, n_cls, [3, 8, 36, 3])
if __name__ == "__main__":
from torchvision.models import resnet
from torchsummaryM import summary
test_model = resnet18(3, 1000)
official_model = resnet.resnet18()
summary(test_model, torch.zeros(1, 3, 224, 224))
summary(official_model, torch.zeros(1, 3, 224, 224))
x = self.classifier(x)
return x
"""
Define VGG16, VGG16_BN, VGG19, VGG19_BN
"""
def vgg16(in_c, n_cls):
return VGG(in_c, n_cls, [2, 2, 3, 3, 3], _bn=False)
def vgg19(in_c, n_cls):
return VGG(in_c, n_cls, [2, 2, 4, 4, 4], _bn=False)
def vgg16_bn(in_c, n_cls):
return VGG(in_c, n_cls, [2, 2, 3, 3, 3], _bn=True)
def vgg19_bn(in_c, n_cls):
return VGG(in_c, n_cls, [2, 2, 4, 4, 4], _bn=True)
if __name__ == "__main__":
from torchsummaryM import summary
from torchvision.models import vgg
official_model = vgg.vgg16()
test_model = vgg16(3, 1000)
inps = torch.zeros(1, 3, 224, 224)
summary(official_model, inps)
summary(test_model, inps)
self.optim.load_state_dict(optim_dict)
try:
self.module.load_state_dict(state_dict)
except AttributeError as e:
self.load_state_dict(state_dict)
self.cfg.start_epoch = epoch
logging.info(f"{ckpt} model loaded.")
logging.info(f"Epoch re-start at {epoch}.\n")
else:
# raise FileNotFoundError("File doesn't exsist. Please check the directory.")
logging.info("File doesn't exsist. Please check the directory.")
logging.info("Start training with initailized model.\n")
if __name__ == "__main__":
from pathlib import Path
from torchsummaryM import summary
class config:
def __init__(self):
# model
self.latent_dim = 16
self.encode_hidden_dims = [64, 128]
self.decode_hidden_dims = [64, 16]
cfg = config()
model = VAE(cfg)
summary(model, torch.zeros((2, 16)))
size, baseline = cv2.getTextSize(text,
cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.5,
thickness=2)
text_w, text_h = size
x, y = left, top
x1y1 = (x, y)
x2y2 = (x + text_w + line_thickness,
y + text_h + line_thickness + baseline)
cv2.rectangle(image_boxes, x1y1, x2y2, bgr, -1)
cv2.putText(image_boxes,
text,
(x + line_thickness, y + 2 * baseline + line_thickness),
cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.4,
color=(255, 255, 255),
thickness=1,
lineType=8)
return image_boxes
if __name__ == "__main__":
from torchsummaryM import summary
from torchvision.models import vgg16_bn
inps = torch.ones(1, 3, 448, 448)
yolo = YOLO_v1(_model='resnet')
summary(yolo, inps)
def main():
args = parser.parse_args()
device = "cuda" if args.cuda else "cpu"
if device == "cuda":
torch.backends.cudnn.benchmark = True
# Set train & valid dataset
dataset = VOC_Dataset(args.img, args.label, debug=args.debug)
valid_ratio = .01
dataset_size = len(dataset)
valid_size = int(valid_ratio * dataset_size)
train_set, valid_set = random_split(dataset, [dataset_size-valid_size, valid_size])
train_loader = DataLoader(train_set, batch_size=args.batch, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=1)
# Define feature_extractor and Yolo model based on
# [You Only Look Once: Unified, Real-Time Object Detection] https://arxiv.org/abs/1506.02640
yolo_model = YOLO_v1(bn=args.bn, _model=args.feature, _device=device).to(device)
# Print summary of the model if needed
if args.summary:
summary(yolo_model, torch.ones(1, 3, 448, 448).to(device))
# Define criterion & optimizer based on
# [You Only Look Once: Unified, Real-Time Object Detection] https://arxiv.org/abs/1506.02640
criterion = YOLO_loss(f_size=yolo_model.S, _gpu=True)
optimizer = torch.optim.SGD(yolo_model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
debug_dir = 'debug'
os.makedirs(debug_dir, exist_ok=True)
for epoch in range(args.start_epoch, args.epochs+1):
yolo_model.train()
cnt = 1
train_loss = 0
train_batch = 0
valid_loss = 0
valid_batch = 0
dataset._train = True
for i, (inps, tars) in enumerate(train_loader):
update_lr(optimizer, epoch, float(i) / float(len(train_loader) - 1), args=args)
lr = get_lr(optimizer)
inps = inps.float().to(device)
tars = tars.float().to(device)
preds = yolo_model(inps)
loss = criterion(preds, tars)
batch_size = inps.size(0)
loss_cur_iter = loss.item()
train_loss += loss_cur_iter * batch_size
train_batch += batch_size
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch [{:3d}/{:3d}] Iter [{:5d}/{:5d}] LR: {:.4f}'.format(
epoch, args.epochs, i, len(train_loader),
lr,
end=''))
print(' Loss: {:.4f} Average Loss: {:.4f}'.format(
loss.item(),
train_loss / float(train_batch)
))
if i % 200 == 0:
yolo_model.eval()
dataset._train = False
if i == 0:
dataset._debug = True
else:
dataset._debug = False
for j, (inps, tars) in enumerate(valid_loader):
# Load data as a batch.
inps = inps.float().to(device)
# Forward to compute validation loss.
with torch.no_grad():
boxes, class_names, probs = yolo_model.detect(image_bgr=inps, image_size=448, valid=True)
img = (inps.squeeze(0).detach().cpu().numpy().transpose(1, 2, 0) * 127.5) + 127.5
plot_img = visualize_boxes(img, boxes, class_names, probs, valid=False)
saveD = os.path.join(debug_dir, 'train_valid', '{:s}'.format(str(epoch).zfill(3)))
os.makedirs(saveD, exist_ok=True)
saveD = os.path.join(saveD, '{:s}'.format(str(i).zfill(5)))
os.makedirs(saveD, exist_ok=True)
cv2.imwrite(os.path.join(saveD, 'debug_{:s}.jpg'.format(str(cnt).zfill(5))), plot_img)
cnt += 1
yolo_model.train()
torch.save(yolo_model.state_dict(), "./yolo_v1.ptr")
Inverted_Residual_Block(160, 320, s=1, t=6, alpha=alpha),
nn.Conv2d(int(320 * alpha), int(alpha * 1280), 1, 1, 0,
bias=False),
nn.BatchNorm2d(int(alpha * 1280)),
nn.ReLU6(True)
]
feature = nn.Sequential(*layers)
return feature
def _build_classifier(self, alpha):
layers = [
nn.AdaptiveAvgPool2d(output_size=(1, 1)),
nn.Flatten(start_dim=1),
nn.Linear(in_features=int(alpha * 1280), out_features=1000),
]
classifier = nn.Sequential(*layers)
return classifier
if __name__ == "__main__":
from torchsummaryM import summary
from torchvision.models.mobilenet import MobileNetV2
model = MobileNetV2()
summary(model, torch.zeros(1, 3, 224, 224))
model = MobileNet_v2(alpha=1.4)
summary(model, torch.zeros(1, 3, 224, 224))
