Spaces:
Sleeping
Sleeping
# coding: utf-8 | |
import torch.nn as nn | |
class C3D(nn.Module): | |
""" | |
nb_classes: nb_classes in classification task, 101 for UCF101 dataset | |
""" | |
def __init__(self, nb_classes): | |
super(C3D, self).__init__() | |
self.conv1 = nn.Conv3d(3, 64, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2)) | |
self.conv2 = nn.Conv3d(64, 128, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.pool2 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2)) | |
self.conv3a = nn.Conv3d(128, 256, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.conv3b = nn.Conv3d(256, 256, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.pool3 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2)) | |
self.conv4a = nn.Conv3d(256, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.conv4b = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.pool4 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2)) | |
self.conv5a = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.conv5b = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1)) | |
self.pool5 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 1, 1)) | |
self.fc6 = nn.Linear(8192, 4096) | |
self.fc7 = nn.Linear(4096, 4096) | |
self.fc8 = nn.Linear(4096, nb_classes) | |
self.dropout = nn.Dropout(p=0.5) | |
self.relu = nn.ReLU() | |
def forward(self, x, feature_layer): | |
h = self.relu(self.conv1(x)) | |
h = self.pool1(h) | |
h = self.relu(self.conv2(h)) | |
h = self.pool2(h) | |
h = self.relu(self.conv3a(h)) | |
h = self.relu(self.conv3b(h)) | |
h = self.pool3(h) | |
h = self.relu(self.conv4a(h)) | |
h = self.relu(self.conv4b(h)) | |
h = self.pool4(h) | |
h = self.relu(self.conv5a(h)) | |
h = self.relu(self.conv5b(h)) | |
h = self.pool5(h) | |
h = h.reshape(-1, 8192) | |
out = h if feature_layer == 5 else None | |
h = self.relu(self.fc6(h)) | |
out = h if feature_layer == 6 and out == None else out | |
h = self.dropout(h) | |
h = self.relu(self.fc7(h)) | |
out = h if feature_layer == 7 and out == None else out | |
h = self.dropout(h) | |
logits = self.fc8(h) | |
return logits, out | |