wchliao / multi-task-image-classification Goto Github PK

Contributors

Stargazers

Forkers

iamweiweishi deepcolin unsterbliche judithyueli imonelittlesheep chenmoshushi kensuncarestream arjunakula manoj652 tongyanjun captain1986 judepark96 wuzhan11 dodoyeon light-- t-mac-curry 1498079723 run-r peternara zongking123 mgsong ankit-iitkgp ryanboldi lijuny andylao-ai 1027jy phanducthien82 deyh2020

def Model(num_classes, num_channels):
    layers = [
        nn.Conv2d(num_channels, 32, kernel_size=3, padding=1),
        nn.BatchNorm2d(32),
        nn.ReLU(),
        nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.Conv2d(64, 64, kernel_size=3, padding=1),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2),
        nn.BatchNorm2d(128),
        nn.ReLU(),
        nn.Conv2d(128, 128, kernel_size=3, padding=1),
        nn.BatchNorm2d(128),
        nn.ReLU(),
    ]

class _Decoder(nn.Module):
    def __init__(self, output_size):
        super(_Decoder, self).__init__()
        self.layers = nn.Sequential(
            # nn.Linear(128*8*8, 512),
            nn.Linear(8 * 112 * 112, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512, output_size)
        )

RuntimeError: CUDA out of memory. Tried to allocate 196.00 MiB (GPU 0; 11.91 GiB total capacity; 11.26 GiB already allocated; 47.06 MiB free; 50.22 MiB cached)